From 2bcb8b084946724aac59367ce7cb1e8471162249 Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Wed, 30 Sep 2026 01:29:25 +0000 Subject: [PATCH 1/9] feat(inference/mlx): on-device self-speculative decoding on Apple Silicon Serve Trida2.0-4B's lossless self-spec (bd_bidir_shift, block 7 / gen 4, strict_truncated) with MLX, on top of mlx-lm's Qwen3.5 model and Metal gated-delta kernel: - canvas forward: full-attention rows 0..N-1 causal, MASK rows see the whole canvas; GDN clean rows read h[t], MASK rows read h[block_end]; nothing is persisted until commit(adv) (KV trim, GDN state of the accepted prefix, conv window from the saved pre-conv rows) - fused canvas GDN Metal kernel (one launch per layer; verify step 49.8 -> 44.5 ms on M3 Pro q8, 1.13x an AR step) - greedy and top-k/top-p speculative sampling, AR baseline, prompt cache with slots and snapshots (GDN layers cannot rewind) - convert (MLX q8/q6/q4), verify, bench, profile tools; uv project - tests on a tiny random model: canvas rows == AR logits, block-end readout, commit == AR state, greedy lossless, sampled distribution == AR M3 Pro 18 GB, q8, greedy: 27 -> 49 tok/s (2.3 tokens per forward). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- inference/mlx/.gitignore | 4 + inference/mlx/.python-version | 1 + inference/mlx/pyproject.toml | 42 ++ inference/mlx/requirements.txt | 5 + inference/mlx/tests/test_tiny.py | 297 +++++++++++ inference/mlx/trida_mlx/__init__.py | 8 + inference/mlx/trida_mlx/bench.py | 147 ++++++ inference/mlx/trida_mlx/convert.py | 85 +++ inference/mlx/trida_mlx/decode.py | 254 +++++++++ inference/mlx/trida_mlx/engine.py | 228 ++++++++ inference/mlx/trida_mlx/kernels.py | 133 +++++ inference/mlx/trida_mlx/model.py | 403 ++++++++++++++ inference/mlx/trida_mlx/profile.py | 166 ++++++ inference/mlx/trida_mlx/verify.py | 92 ++++ inference/mlx/uv.lock | 783 ++++++++++++++++++++++++++++ 15 files changed, 2648 insertions(+) create mode 100644 inference/mlx/.gitignore create mode 100644 inference/mlx/.python-version create mode 100644 inference/mlx/pyproject.toml create mode 100644 inference/mlx/requirements.txt create mode 100644 inference/mlx/tests/test_tiny.py create mode 100644 inference/mlx/trida_mlx/__init__.py create mode 100644 inference/mlx/trida_mlx/bench.py create mode 100644 inference/mlx/trida_mlx/convert.py create mode 100644 inference/mlx/trida_mlx/decode.py create mode 100644 inference/mlx/trida_mlx/engine.py create mode 100644 inference/mlx/trida_mlx/kernels.py create mode 100644 inference/mlx/trida_mlx/model.py create mode 100644 inference/mlx/trida_mlx/profile.py create mode 100644 inference/mlx/trida_mlx/verify.py create mode 100644 inference/mlx/uv.lock diff --git a/inference/mlx/.gitignore b/inference/mlx/.gitignore new file mode 100644 index 0000000..cdedea9 --- /dev/null +++ b/inference/mlx/.gitignore @@ -0,0 +1,4 @@ +# local benchmark / profiling outputs and request logs +bench*.json +reqlog/ +*.egg-info/ diff --git a/inference/mlx/.python-version b/inference/mlx/.python-version new file mode 100644 index 0000000..e4fba21 --- /dev/null +++ b/inference/mlx/.python-version @@ -0,0 +1 @@ +3.12 diff --git a/inference/mlx/pyproject.toml b/inference/mlx/pyproject.toml new file mode 100644 index 0000000..2b6714e --- /dev/null +++ b/inference/mlx/pyproject.toml @@ -0,0 +1,42 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "trida-mlx" +version = "0.1.0" +description = "On-device (Apple Silicon / MLX) self-speculative decoding, OpenAI server and agent for Trida2.0-4B" +readme = "README.md" +requires-python = ">=3.10,<3.14" +license = { text = "Apache-2.0" } +authors = [{ name = "Trillion Labs" }] +dependencies = [ + # Metal build on macOS; the CPU build elsewhere (Linux CI / offline tests only). + "mlx>=0.29; sys_platform == 'darwin'", + "mlx[cpu]>=0.29; sys_platform == 'linux'", + "mlx-lm>=0.31", # Qwen3.5 gated-delta model + Metal kernel we build on + "transformers>=4.57", # tokenizer + chat template (tools / enable_thinking) + "huggingface-hub", +] + +[project.scripts] +trida-mlx-server = "trida_mlx.server:main" +trida-mlx-agent = "trida_mlx.agent:main" +trida-mlx-convert = "trida_mlx.convert:main" +trida-mlx-verify = "trida_mlx.verify:main" +trida-mlx-bench = "trida_mlx.bench:main" +trida-mlx-profile = "trida_mlx.profile:main" + +[dependency-groups] +dev = ["pytest>=8"] + +[tool.setuptools] +packages = ["trida_mlx"] + +[tool.uv] +# resolve only for the platforms this package runs on +environments = ["sys_platform == 'darwin' and platform_machine == 'arm64'", "sys_platform == 'linux'"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = [".", "tests"] diff --git a/inference/mlx/requirements.txt b/inference/mlx/requirements.txt new file mode 100644 index 0000000..2830b61 --- /dev/null +++ b/inference/mlx/requirements.txt @@ -0,0 +1,5 @@ +# On-device (Apple Silicon) self-speculative serving for Trida2.0-4B. +mlx>=0.29 +mlx-lm>=0.31 # provides the Qwen3.5 gated-delta model + Metal kernel we build on +transformers>=4.57 # tokenizer + chat template (tools / enable_thinking) +huggingface_hub diff --git a/inference/mlx/tests/test_tiny.py b/inference/mlx/tests/test_tiny.py new file mode 100644 index 0000000..7c9f01e --- /dev/null +++ b/inference/mlx/tests/test_tiny.py @@ -0,0 +1,297 @@ +"""Correctness tests on a tiny random Qwen3.5-hybrid model (CPU is fine; no weights needed). + + python -m pytest inference/mlx/tests -q # or: python inference/mlx/tests/test_tiny.py + +What is pinned: + * chunked prefill == one-shot prefill + * canvas clean rows (0..N-1) reproduce the AR logits exactly (the lossless property) + * canvas MASK rows read the *block-end* GDN state (causal_mode=2 semantics) + * commit(adv) leaves the cache identical to having AR-decoded those adv tokens + * greedy self-spec output == greedy AR output, token for token, over long generations + * a snapshot/restore round trip (prompt cache) reproduces a fresh prefill +""" +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from trida_mlx.decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate +from trida_mlx.model import GDNState, TridaRuntime, build_model +from mlx_lm.models.gated_delta import compute_g + +MASK = 299 +TINY = dict( + model_type="qwen3_5_text", hidden_size=64, intermediate_size=128, num_hidden_layers=8, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, rms_norm_eps=1e-6, vocab_size=300, + linear_num_value_heads=4, linear_num_key_heads=2, linear_key_head_dim=32, linear_value_head_dim=16, + linear_conv_kernel_dim=4, tie_word_embeddings=True, full_attention_interval=4, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.25, + "mrope_section": [1, 1, 0]}, +) + + +def tiny_runtime(seed=0, sharp=4.0): + mx.random.seed(seed) + model = build_model(TINY) + # sharpen the output distribution so greedy decoding has structure (and fewer near-ties) + emb = model.language_model.model.embed_tokens + emb.weight = emb.weight * sharp + mx.eval(model.parameters()) + return TridaRuntime(model, prefill_chunk=512) + + +def rand_prompt(n, seed=1): + rng = np.random.default_rng(seed) + return rng.integers(0, 290, size=n).tolist() + + +def close(a, b, tol=1e-4): + return float(mx.max(mx.abs(a - b)).item()) < tol + + +def test_chunked_prefill(): + rt = tiny_runtime() + p = rand_prompt(37) + c1 = rt.make_cache(); l1 = rt.prefill(c1, p) + rt2 = tiny_runtime(); rt2.prefill_chunk = 5 + c2 = rt2.make_cache(); l2 = rt2.prefill(c2, p) + assert close(l1, l2) + + +def _ar_logits(rt, prompt, cont): + c = rt.make_cache() + out = [rt.prefill(c, prompt)] + for t in cont: + out.append(rt.ar_step(c, t)) + return out, c + + +def test_canvas_clean_rows_are_ar_logits(): + rt = tiny_runtime(); n = 4 + prompt = rand_prompt(20) + cont = rand_prompt(n, seed=7) # p, s0, s1, s2 + ar, _ = _ar_logits(rt, prompt, cont) + c = rt.make_cache(); rt.prefill(c, prompt) + L = rt.canvas(c, cont + [MASK] * (n - 1), n) + for i in range(n): # row i == AR logits after cont[:i+1] + assert close(L[i], ar[i + 1]), f"row {i}" + + +def test_commit_matches_ar_state(): + n = 4 + for adv in range(1, n + 1): + rt = tiny_runtime() + prompt = rand_prompt(15) + cont = rand_prompt(n, seed=3) + c = rt.make_cache(); rt.prefill(c, prompt) + rt.canvas(c, cont + [MASK] * (n - 1), n) + rt.commit(c, adv) + nxt = 42 + got = rt.ar_step(c, nxt) + ar, _ = _ar_logits(rt, prompt, cont[:adv] + [nxt]) + assert close(got, ar[-1]), f"adv={adv}" + assert c.tokens == prompt + cont[:adv] + [nxt] + + +def test_mask_rows_read_block_end_state(): + """Independent re-derivation of one GDN layer's canvas output with a per-token loop.""" + rt = tiny_runtime(); n = 4 + layer = next(l for l in rt.backbone.layers if l.is_linear) + lin = layer.linear_attn + x = mx.random.normal((1, 2 * n - 1, TINY["hidden_size"])) + st = GDNState(conv=mx.random.normal((1, 3, lin.conv_dim)), + ssm=mx.random.normal((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim)) * 0.1) + out, _, _ = rt._gdn(lin, x, st, n) + # manual: conv -> q,k,v -> step recurrence, record states + qkv = lin.in_proj_qkv(x); z = lin.in_proj_z(x).reshape(1, 2 * n - 1, lin.num_v_heads, lin.head_v_dim) + ci = mx.concatenate([st.conv, qkv], axis=1) + co = nn.silu(lin.conv1d(ci)) + q, k, v = mx.split(co, [lin.key_dim, 2 * lin.key_dim], -1) + T = 2 * n - 1 + q = q.reshape(1, T, lin.num_k_heads, -1); k = k.reshape(1, T, lin.num_k_heads, -1) + v = v.reshape(1, T, lin.num_v_heads, -1) + s = k.shape[-1] ** -0.5 + q = s * s * mx.fast.rms_norm(q, None, 1e-6); k = s * mx.fast.rms_norm(k, None, 1e-6) + rep = lin.num_v_heads // lin.num_k_heads + q = mx.repeat(q, rep, 2); k = mx.repeat(k, rep, 2) + beta = mx.sigmoid(lin.in_proj_b(x)); g = compute_g(lin.A_log, lin.in_proj_a(x), lin.dt_bias) + S = st.ssm; ys = [] + for t in range(T): + S = S * g[:, t][..., None, None] + kv = (S * k[:, t][:, :, None, :]).sum(-1) + S = S + k[:, t][:, :, None, :] * ((v[:, t] - kv) * beta[:, t][..., None])[..., None] + ys.append((S * q[:, t][:, :, None, :]).sum(-1)) + S_end = S + for t in range(n, T): + ys[t] = (S_end * q[:, t][:, :, None, :]).sum(-1) + y = mx.stack(ys, axis=1) + ref = lin.out_proj(lin.norm(y, z).reshape(1, T, -1)) + assert close(out, ref, 1e-4) + + +def test_fused_gdn_matches_unfused(): + """Metal-only: the fused canvas GDN kernel == two mlx-lm kernel calls + block-end readout.""" + from trida_mlx.kernels import fused_available + if not fused_available(): + print(" (skipped: no Metal)") + return + rt = tiny_runtime(); n = 4 + lin = next(l for l in rt.backbone.layers if l.is_linear).linear_attn + x = mx.random.normal((1, 2 * n - 1, TINY["hidden_size"])) + st = GDNState(conv=mx.random.normal((1, 3, lin.conv_dim)), + ssm=mx.random.normal((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim)) * 0.1) + rt.fused_gdn = True + o1, _, r1 = rt._gdn(lin, x, st, n) + rt.fused_gdn = False + o2, _, r2 = rt._gdn(lin, x, st, n) + assert close(o1, o2, 1e-5) and close(r1.s_clean, r2.s_clean, 1e-6) + # clean rows must be bit-identical to the unfused (= AR) kernel + assert bool(mx.array_equal(o1[:, :n], o2[:, :n]).item()) + + +def _collect(gen): + out = [] + for chunk in gen: + out.extend(chunk) + return out + + +def test_greedy_selfspec_is_lossless(): + sp = SamplingParams(temperature=0.0) + for seed in range(3): + rt = tiny_runtime(seed=seed) + prompt = rand_prompt(25, seed=seed + 10) + ar = _collect(ar_generate(rt, rt.make_cache(), prompt, max_new_tokens=120, eos_ids=set(), sp=sp)) + from trida_mlx.decode import DecodeStats + st = DecodeStats() + ss = _collect(selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=120, eos_ids=set(), + sp=sp, mask_id=MASK, n=4, stats=st)) + assert ss == ar, f"seed {seed}: first diff at {next(i for i,(a,b) in enumerate(zip(ss,ar)) if a!=b)}" + assert st.forwards <= len(ar) + + +def test_greedy_lossless_with_oracle_drafts(): + """Random models rarely accept drafts, so also drive the accept path: patch the canvas + so MASK rows 'predict' the true AR continuation and check output + acceptance.""" + sp = SamplingParams(temperature=0.0) + rt = tiny_runtime(seed=5) + prompt = rand_prompt(18, seed=4) + ar = _collect(ar_generate(rt, rt.make_cache(), prompt, max_new_tokens=60, eos_ids=set(), sp=sp)) + full = prompt + ar + orig = rt.canvas + + # replace the MASK draft rows with a one-hot on the true AR continuation + def patched(cache, toks, n, rows=None): + base = cache.length + L = orig(cache, toks, n, rows=rows) + rows = [] + V = L.shape[-1] + for i in range(L.shape[0]): + want = base + i + 1 + if i >= 1 and want < len(full) and (toks[i] == MASK): + rows.append(mx.where(mx.arange(V) == full[want], 100.0, 0.0)) + else: + rows.append(L[i]) + return mx.stack(rows) + + rt.canvas = patched + from trida_mlx.decode import DecodeStats + st = DecodeStats() + ss = _collect(selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=60, eos_ids=set(), + sp=sp, mask_id=MASK, n=4, stats=st)) + assert ss == ar + assert st.accept_hist[-1] > 0 and st.tokens_per_forward > 2.5, st.as_dict() + + +def test_snapshot_restore(): + rt = tiny_runtime() + p = rand_prompt(30) + c = rt.make_cache(); rt.prefill(c, p[:20]) + snap = c.snapshot() + rt.prefill(c, p[20:]) + for t in rand_prompt(9, seed=9): + rt.ar_step(c, t) + c.restore(snap) + l1 = rt.prefill(c, p[20:]) + l2 = rt.prefill(rt.make_cache(), p) + assert close(l1, l2) + + +def test_sampling_runs(): + sp = SamplingParams(temperature=1.0, top_k=20, top_p=0.9, seed=0) + rt = tiny_runtime() + out = _collect(selfspec_generate(rt, rt.make_cache(), rand_prompt(12), max_new_tokens=40, + eos_ids=set(), sp=sp, mask_id=MASK, n=4)) + assert len(out) == 40 + + +def test_truncated_probs_matches_reference(): + from trida_mlx.decode import truncated_probs + + def ref(x, t, k, p): # numpy port of the SGLang reference _sampling_probs + x = x / t + kth = np.sort(x, axis=-1)[:, -k][:, None] + x = np.where(x < kth, -np.inf, x) + order = np.argsort(-x, axis=-1, kind="stable") + sx = np.take_along_axis(x, order, -1) + e = np.exp(sx - sx.max(-1, keepdims=True)); pr = e / e.sum(-1, keepdims=True) + sx = np.where(np.cumsum(pr, -1) - pr >= p, -np.inf, sx) + out = np.full_like(x, -np.inf); np.put_along_axis(out, order, sx, -1) + e = np.exp(out - out.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True) + + rng = np.random.default_rng(0) + x = rng.normal(size=(7, 1000)).astype(np.float32) * 3 + for t, k, p in [(1.0, 50, 0.95), (0.7, 20, 0.8), (1.0, 5, 1.0)]: + got = np.array(truncated_probs(mx.array(x), SamplingParams(temperature=t, top_k=k, top_p=p))) + assert np.allclose(got, ref(x, t, k, p), atol=1e-5), (t, k, p) + + +def test_speculative_sampling_preserves_distribution(): + """Sampled self-spec must reproduce the AR sampling distribution. Random models almost + never accept drafts, so the MASK rows are patched to reuse a neighbouring row's logits: + any draft q is valid for speculative sampling, and this one gets accepted often.""" + from collections import Counter + rt = tiny_runtime(seed=2, sharp=1.0) + prompt = rand_prompt(10, seed=2) + sp = SamplingParams(temperature=1.0, top_k=4, top_p=0.9) + orig = rt.canvas + + def patched(cache, toks, n, rows=None): + L = orig(cache, toks, n, rows=rows) + return mx.stack([L[i] if (i == 0 or toks[i] != MASK) else L[i - 1] for i in range(L.shape[0])]) + + def hist(fn, reps, seed0, **kw): + c, st_all = Counter(), DecodeStats() + st_all.accept_hist = [0] * 4 + for i in range(reps): + mx.random.seed(seed0 + i) + st = DecodeStats() + out = _collect(fn(rt, rt.make_cache(), prompt, max_new_tokens=4, eos_ids=set(), sp=sp, stats=st, **kw)) + c[out[3]] += 1 + if st.accept_hist: + st_all.accept_hist = [x + y for x, y in zip(st_all.accept_hist, st.accept_hist)] + return c, st_all + + R = 600 + a, _ = hist(ar_generate, R, 5000) + a2, _ = hist(ar_generate, R, 90000) + rt.canvas = patched + b, st = hist(selfspec_generate, R, 7000, mask_id=MASK, n=4) + tv = lambda x, y: 0.5 * sum(abs(x[k] - y[k]) for k in set(x) | set(y)) / R + assert sum(st.accept_hist[1:]) > R // 4, st.accept_hist # the accept path really ran + assert tv(a, b) < max(0.08, 2.0 * tv(a, a2)), (tv(a, b), tv(a, a2)) + + +if __name__ == "__main__": + fails = 0 + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + try: + fn(); print(f" ok {name}") + except Exception as e: # noqa: BLE001 + fails += 1; print(f" FAIL {name}: {type(e).__name__}: {e}") + sys.exit(1 if fails else 0) diff --git a/inference/mlx/trida_mlx/__init__.py b/inference/mlx/trida_mlx/__init__.py new file mode 100644 index 0000000..9bc8e09 --- /dev/null +++ b/inference/mlx/trida_mlx/__init__.py @@ -0,0 +1,8 @@ +"""trida_mlx — on-device (Apple Silicon / MLX) self-speculative decoding for Trida2.0-4B.""" +from .model import DEFAULT_MODEL, SeqCache, TridaRuntime, load_model, mask_id_of +from .decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate + +__all__ = [ + "DEFAULT_MODEL", "SeqCache", "TridaRuntime", "load_model", "mask_id_of", + "DecodeStats", "SamplingParams", "ar_generate", "selfspec_generate", +] diff --git a/inference/mlx/trida_mlx/bench.py b/inference/mlx/trida_mlx/bench.py new file mode 100644 index 0000000..08d75f4 --- /dev/null +++ b/inference/mlx/trida_mlx/bench.py @@ -0,0 +1,147 @@ +"""Benchmark causal vs self-spec on device, and check that greedy self-spec is lossless. + + python -m trida_mlx.bench --model ./Trida2.0-4B-mlx-q8 # built-in prompts + python -m trida_mlx.bench --model ./Trida2.0-4B-mlx-q8 --temperature 1.0 --repeats 3 + python -m trida_mlx.bench --model ./Trida2.0-4B-mlx-q8 --prompts my_prompts.jsonl --out results.json + +Reports per prompt and in aggregate: decode tok/s, tokens per forward, accept histogram, +prefill time, and (greedy) whether self-spec output == causal output token for token. +Prompts file: JSONL with {"messages": [...]} or {"prompt": "..."} (+ optional "tools"). +""" +from __future__ import annotations + +import argparse +import json +import platform +import statistics +import subprocess +import time + +import mlx.core as mx + +from .decode import DecodeStats, SamplingParams +from .engine import Engine + +BUILTIN = [ + {"name": "gsm8k", "messages": [{"role": "user", "content": + "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. " + "How many clips did Natalia sell altogether in April and May? Reason step by step and put the " + "final answer after '####'."}]}, + {"name": "code", "messages": [{"role": "user", "content": + "Write a Python function that returns the n-th Fibonacci number iteratively, with a docstring " + "and two doctests."}]}, + {"name": "tool_call", "tools": [{"type": "function", "function": { + "name": "get_weather", "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}, + "unit": {"type": "string", "enum": ["c", "f"]}}, + "required": ["city"]}}}], + "messages": [{"role": "user", "content": "What's the weather in Seoul right now, in celsius?"}]}, + {"name": "korean", "messages": [{"role": "user", "content": + "온디바이스 LLM의 장점과 단점을 각각 세 가지씩 설명해 주세요."}]}, +] + + +def _hw() -> str: + try: + chip = subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip() + mem = int(subprocess.run(["sysctl", "-n", "hw.memsize"], capture_output=True, text=True).stdout.strip() or 0) + return f"{chip}, {mem / 2**30:.0f}GB" + except Exception: + return platform.platform() + + +def run_one(eng: Engine, prompt_ids, mode, sp, max_tokens): + eng.reset_cache() # cold prompt for a fair prefill number + st = DecodeStats() + ids = [] + t = time.perf_counter() + for chunk in eng.generate_ids(prompt_ids, max_tokens=max_tokens, sampling=sp, mode=mode, stats=st): + ids.extend(chunk) + wall = time.perf_counter() - t + d = st.as_dict() + d["wall_s"] = round(wall, 3) + return ids, d + + +def main(argv=None): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") + ap.add_argument("--prompts", default=None) + ap.add_argument("--modes", default="causal,self-spec") + ap.add_argument("--gen-block", type=int, default=4) + ap.add_argument("--max-tokens", type=int, default=512) + ap.add_argument("--temperature", type=float, default=0.0) + ap.add_argument("--top-k", type=int, default=50) + ap.add_argument("--top-p", type=float, default=0.95) + ap.add_argument("--repeats", type=int, default=1) + ap.add_argument("--no-think", action="store_true") + ap.add_argument("--out", default=None) + ap.add_argument("--no-fused-gdn", action="store_true", help="A/B: unfused canvas GDN (2 kernels + readout)") + ap.add_argument("--skip-rows", action="store_true", help="A/B: LM head on rows 0..N-1 only at cold start") + a = ap.parse_args(argv) + + eng = Engine(a.model, gen_block=a.gen_block, fused_gdn=not a.no_fused_gdn) + eng.rt.skip_rows = a.skip_rows + items = BUILTIN + if a.prompts: + items = [json.loads(l) for l in open(a.prompts) if l.strip()] + modes = [m.strip() for m in a.modes.split(",") if m.strip()] + print(f"hardware: {_hw()} model: {a.model} load {eng.load_s:.1f}s mask_id={eng.mask_id}") + print("warming up ...", flush=True) + for m in modes: + for _ in eng.generate_ids(eng.encode("<|im_start|>user\nhello<|im_end|>\n<|im_start|>assistant\n"), + max_tokens=16, sampling=SamplingParams(temperature=0.0), mode=m): + pass + results = [] + for it in items: + text = it.get("prompt") or eng.render(it["messages"], tools=it.get("tools"), + enable_thinking=not a.no_think) + pids = eng.encode(text) + row = {"name": it.get("name", f"p{len(results)}"), "prompt_tokens": len(pids), "runs": {}} + outs = {} + for m in modes: + runs = [] + for r in range(a.repeats): + sp = SamplingParams(temperature=a.temperature, top_k=a.top_k, top_p=a.top_p, + seed=None if a.temperature <= 0 else r) + ids, d = run_one(eng, pids, m, sp, a.max_tokens) + runs.append(d) + outs.setdefault(m, ids) + row["runs"][m] = runs + if a.temperature <= 0 and "causal" in outs and "self-spec" in outs: + A, B = outs["causal"], outs["self-spec"] + n = min(len(A), len(B)) + first = next((i for i in range(n) if A[i] != B[i]), None if len(A) == len(B) else n) + row["identical"] = first is None + row["first_diff"] = first + row["sample"] = eng.tokenizer.decode(outs[modes[-1]][:160]) + results.append(row) + line = f"[{row['name']:>9}] prompt={len(pids):5d}" + for m in modes: + tps = statistics.mean(r["decode_tok_s"] for r in row["runs"][m]) + tpf = statistics.mean(r["tokens_per_forward"] for r in row["runs"][m]) + nt = statistics.mean(r["new_tokens"] for r in row["runs"][m]) + line += f" | {m}: {tps:6.1f} tok/s {tpf:4.2f} tok/fwd {nt:5.0f} tok" + if "identical" in row: + line += f" | lossless={'YES' if row['identical'] else 'NO@' + str(row['first_diff'])}" + print(line, flush=True) + + agg = {} + for m in modes: + tok = sum(r["new_tokens"] for row in results for r in row["runs"][m]) + sec = sum(r["decode_s"] for row in results for r in row["runs"][m]) + fw = sum(r["forwards"] for row in results for r in row["runs"][m]) + agg[m] = {"decode_tok_s": round(tok / sec, 2) if sec else 0, "tokens_per_forward": round(tok / fw, 3) if fw else 0} + if "causal" in agg and "self-spec" in agg and agg["causal"]["decode_tok_s"]: + agg["speedup"] = round(agg["self-spec"]["decode_tok_s"] / agg["causal"]["decode_tok_s"], 3) + print("aggregate:", json.dumps(agg)) + if a.out: + with open(a.out, "w") as f: + json.dump({"hardware": _hw(), "model": a.model, "args": vars(a), "aggregate": agg, + "results": results, "peak_memory_gb": round(mx.get_peak_memory() / 1e9, 2)}, f, + indent=2, ensure_ascii=False) + print(f"-> {a.out}") + + +if __name__ == "__main__": + main() diff --git a/inference/mlx/trida_mlx/convert.py b/inference/mlx/trida_mlx/convert.py new file mode 100644 index 0000000..abcaf2f --- /dev/null +++ b/inference/mlx/trida_mlx/convert.py @@ -0,0 +1,85 @@ +"""Convert / quantize a Trida checkpoint to an MLX directory. + + python -m trida_mlx.convert --model trillionlabs/Trida2.0-4B --out ./Trida2.0-4B-mlx-q8 --bits 8 + python -m trida_mlx.convert --model trillionlabs/Trida2.0-4B --out ./Trida2.0-4B-mlx-bf16 + +The output keeps the tokenizer, chat template, generation config and block_diffusion.json, +and writes a config with ``model_type: qwen3_5`` + ``quantization``, so it also loads with +stock ``mlx_lm.load`` (useful as an AR baseline: ``mlx_lm.server``). + +Quantization notes (from the ocr-monorepo MLX study): 8-bit and 6-bit kept greedy outputs +identical to bf16 there, 4-bit changed them. Self-spec stays lossless *relative to the +quantized AR model* at any bit width; check quality with ``bench.py --compare-ref``. +""" +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten + +from .model import _text_config, load_model + +CARRY = ["tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "generation_config.json", + "block_diffusion.json", "special_tokens_map.json", "vocab.json", "merges.txt", "added_tokens.json"] + + +def main(argv=None): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") + ap.add_argument("--out", required=True) + ap.add_argument("--bits", type=int, default=None, choices=[3, 4, 5, 6, 8], help="omit for bf16") + ap.add_argument("--group-size", type=int, default=64) + ap.add_argument("--keep-embed", action="store_true", + help="keep embed_tokens / (tied) LM head in bf16 (bigger, safer quality)") + ap.add_argument("--shard-gb", type=float, default=4.0) + a = ap.parse_args(argv) + + model, cfg, src = load_model(a.model) + text = _text_config(cfg) + out_cfg = dict(text) + out_cfg["model_type"] = "qwen3_5" + out_cfg["architectures"] = ["Qwen3_5ForCausalLM"] + if a.bits: + def pred(path, m): + if not hasattr(m, "to_quantized"): + return False + if a.keep_embed and ("embed_tokens" in path or "lm_head" in path): + return False + w = getattr(m, "weight", None) + return w is not None and w.shape[-1] % a.group_size == 0 + nn.quantize(model, group_size=a.group_size, bits=a.bits, class_predicate=pred) + out_cfg["quantization"] = {"group_size": a.group_size, "bits": a.bits} + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + weights = dict(tree_flatten(model.parameters())) + shards, cur, cur_bytes = [], {}, 0 + limit = int(a.shard_gb * (1 << 30)) + for k, v in weights.items(): + if cur and cur_bytes + v.nbytes > limit: + shards.append(cur); cur, cur_bytes = {}, 0 + cur[k] = v; cur_bytes += v.nbytes + shards.append(cur) + index = {"metadata": {"total_size": sum(v.nbytes for v in weights.values())}, "weight_map": {}} + for i, sh in enumerate(shards): + name = "model.safetensors" if len(shards) == 1 else f"model-{i + 1:05d}-of-{len(shards):05d}.safetensors" + mx.save_safetensors(str(out / name), sh, metadata={"format": "mlx"}) + for k in sh: + index["weight_map"][k] = name + if len(shards) > 1: + (out / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + (out / "config.json").write_text(json.dumps(out_cfg, indent=2)) + for name in CARRY: + if (src / name).exists(): + shutil.copy2(src / name, out / name) + gb = index["metadata"]["total_size"] / 1e9 + print(f"wrote {out} ({gb:.2f} GB, {len(shards)} shard(s), " + f"{'bf16' if not a.bits else f'q{a.bits} g{a.group_size}'})") + + +if __name__ == "__main__": + main() diff --git a/inference/mlx/trida_mlx/decode.py b/inference/mlx/trida_mlx/decode.py new file mode 100644 index 0000000..178a271 --- /dev/null +++ b/inference/mlx/trida_mlx/decode.py @@ -0,0 +1,254 @@ +"""Decoders: autoregressive baseline and self-speculative (AR-Trust) decoding. + +Self-spec follows the SGLang reference ``HybridDiffusionSelfSpec`` (variant +``bd_bidir_shift``, ``block_size = 2N-1``, ``gen_block_size = N``): + + cold start : [t0, MASK x (2N-2)] -> clean = row 0, specs = rows 1..N-1 + verify : [pending, spec_0..spec_{N-2}, MASK x (N-1)] + spec_i is verified against row i (exact AR logits); + all accepted -> clean = row N-1, new specs = rows N..2N-2 + reject at i -> emit spec_0..spec_{i-1} + correction; next step is a + cold start from the correction token. + +Greedy (temperature 0) is lossless: the output equals ``ar_generate`` token for token. +With temperature > 0 the ``strict_truncated`` policy is used: drafts are sampled from the +same top-k/top-p truncated distribution the verifier uses, accepted with prob +``min(1, p/q)`` and corrected from ``norm(max(0, p - q))`` — standard speculative +sampling, so the output *distribution* matches AR sampling. +""" +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Iterator, Optional + +import mlx.core as mx + +from .model import SeqCache, TridaRuntime + + +# ----------------------------------------------------------------------------- sampling +@dataclass +class SamplingParams: + temperature: float = 1.0 + top_k: int = 50 + top_p: float = 0.95 + seed: Optional[int] = None + + @property + def greedy(self) -> bool: + return self.temperature <= 0 + + +def truncated_probs(logits: mx.array, sp: SamplingParams) -> mx.array: + """Top-k / top-p truncated sampling distribution, float32 [..., V]. + + Same distribution as the reference ``_sampling_probs`` (top-k, then top-p keeping the + token that crosses ``top_p``), but with top-k set the nucleus is computed on the k + survivors only, so there is no full-vocabulary sort (V = 248k).""" + x = logits.astype(mx.float32) + if sp.temperature != 1.0: + x = x / sp.temperature + V = x.shape[-1] + if 0 < sp.top_k < V: + top = mx.sort(mx.topk(x, sp.top_k, axis=-1), axis=-1)[..., ::-1] # [.., k] descending + kth = top[..., -1:] + x = mx.where(x < kth, -mx.inf, x) + if sp.top_p < 1.0: + ps = mx.softmax(top, axis=-1) + cum = mx.cumsum(ps, axis=-1) + keep = (cum - ps) < sp.top_p # a prefix of the sorted survivors + n_keep = mx.maximum(keep.sum(axis=-1, keepdims=True), 1) + thresh = mx.take_along_axis(top, n_keep - 1, axis=-1) + x = mx.where(x < thresh, -mx.inf, x) + elif sp.top_p < 1.0: + order = mx.argsort(-x, axis=-1) + sx = mx.take_along_axis(x, order, axis=-1) + sp_ = mx.softmax(sx, axis=-1) + cum = mx.cumsum(sp_, axis=-1) + sx = mx.where((cum - sp_) >= sp.top_p, -mx.inf, sx) + x = mx.take_along_axis(sx, mx.argsort(order, axis=-1), axis=-1) + return mx.softmax(x, axis=-1) + + +def sample_from_probs(probs: mx.array) -> mx.array: + return mx.random.categorical(mx.log(probs), axis=-1) + + +@dataclass +class DecodeStats: + mode: str = "" + prompt_tokens: int = 0 + reused_tokens: int = 0 + new_tokens: int = 0 + forwards: int = 0 + cold_forwards: int = 0 + accept_hist: list = field(default_factory=list) + prefill_s: float = 0.0 + decode_s: float = 0.0 + finish_reason: str = "length" + + @property + def tokens_per_forward(self) -> float: + return self.new_tokens / self.forwards if self.forwards else 0.0 + + @property + def decode_tps(self) -> float: + return self.new_tokens / self.decode_s if self.decode_s else 0.0 + + def as_dict(self) -> dict: + return { + "mode": self.mode, "prompt_tokens": self.prompt_tokens, "reused_tokens": self.reused_tokens, + "new_tokens": self.new_tokens, "forwards": self.forwards, "cold_forwards": self.cold_forwards, + "tokens_per_forward": round(self.tokens_per_forward, 3), "accept_hist": self.accept_hist, + "prefill_s": round(self.prefill_s, 4), "decode_s": round(self.decode_s, 4), + "decode_tok_s": round(self.decode_tps, 2), "finish_reason": self.finish_reason, + } + + +class _Emitter: + """Applies max_new_tokens / EOS to emitted chunks and records stats.""" + + def __init__(self, stats: DecodeStats, max_new: int, eos: set): + self.stats, self.max_new, self.eos, self.done = stats, max_new, eos, False + + def __call__(self, toks: list) -> list: + out = [] + for t in toks: + if t in self.eos: + self.done = True + self.stats.finish_reason = "stop" + break + out.append(t) + if self.stats.new_tokens + len(out) >= self.max_new: + self.done = True + self.stats.finish_reason = "length" + break + self.stats.new_tokens += len(out) + return out + + +def _prefill(rt: TridaRuntime, cache: SeqCache, prompt: list, stats: DecodeStats): + t = time.perf_counter() + todo = prompt[cache.length:] + stats.prompt_tokens, stats.reused_tokens = len(prompt), cache.length + if not todo: # whole prompt cached: re-feed the last token to get its logits + raise ValueError("prompt fully cached; callers must leave at least one token to prefill") + logits = rt.prefill(cache, todo) + mx.eval(logits) + stats.prefill_s = time.perf_counter() - t + return logits + + +# ----------------------------------------------------------------------------- AR +def ar_generate(rt: TridaRuntime, cache: SeqCache, prompt: list, *, max_new_tokens: int, + eos_ids: set, sp: SamplingParams, stats: Optional[DecodeStats] = None) -> Iterator[list]: + stats = stats or DecodeStats() + stats.mode = "causal" + if sp.seed is not None: + mx.random.seed(sp.seed) + logits = _prefill(rt, cache, prompt, stats) + emit = _Emitter(stats, max_new_tokens, eos_ids) + t = time.perf_counter() + while True: + tok = mx.argmax(logits) if sp.greedy else sample_from_probs(truncated_probs(logits, sp)) + tok = int(tok.item()) + out = emit([tok]) + if out: + yield out + if emit.done: + break + logits = rt.ar_step(cache, tok) + stats.forwards += 1 + stats.decode_s = time.perf_counter() - t + + +# ----------------------------------------------------------------------------- self-spec +def selfspec_generate(rt: TridaRuntime, cache: SeqCache, prompt: list, *, max_new_tokens: int, + eos_ids: set, sp: SamplingParams, mask_id: int, n: int = 4, + stats: Optional[DecodeStats] = None) -> Iterator[list]: + stats = stats or DecodeStats() + stats.mode = f"self-spec(N={n})" + stats.accept_hist = [0] * n # accepted specs per verify round: 0..N-1 + if sp.seed is not None: + mx.random.seed(sp.seed) + blk, K = 2 * n - 1, n - 1 + logits = _prefill(rt, cache, prompt, stats) + emit = _Emitter(stats, max_new_tokens, eos_ids) + t_start = time.perf_counter() + + if sp.greedy: + t0 = int(mx.argmax(logits).item()) + else: + t0 = int(sample_from_probs(truncated_probs(logits, sp)).item()) + t0_emitted = False + pending, specs, spec_q = None, None, None # spec_q: [K, V] draft distributions + + while not emit.done: + if specs is not None: # ---- verify round + canvas = [pending] + specs + [mask_id] * (n - 1) + L = rt.canvas(cache, canvas, n) + stats.forwards += 1 + if sp.greedy: + am = mx.argmax(L, axis=-1) + vals = am.tolist() + acc = 0 + while acc < K and vals[acc] == specs[acc]: + acc += 1 + corr = vals[acc] if acc < K else None + clean, new_specs = vals[n - 1], vals[n:blk] + new_q = None + else: + P = truncated_probs(L, sp) # [blk, V] + sv = mx.array(specs, dtype=mx.int32) + p = mx.take_along_axis(P[:K], sv[:, None], axis=-1)[:, 0] + q = mx.take_along_axis(spec_q, sv[:, None], axis=-1)[:, 0] + ratio = mx.where(q > 0, p / mx.maximum(q, 1e-30), 0.0) + u = mx.random.uniform(shape=(K,)) + accepted = (ratio >= 1.0) | (u < ratio) + resid = mx.maximum(P[:K] - spec_q, 0.0) + rs = resid.sum(axis=-1, keepdims=True) + resid = mx.where(rs > 0, resid / mx.maximum(rs, 1e-30), P[:K]) + corr_all = sample_from_probs(resid) + samp = sample_from_probs(P[n - 1 : blk]) # clean + new specs + packed = mx.concatenate([accepted.astype(mx.int32), corr_all.astype(mx.int32), + samp.astype(mx.int32)]) + vals = packed.tolist() + acc_l, corr_l, samp_l = vals[:K], vals[K : 2 * K], vals[2 * K :] + acc = 0 + while acc < K and acc_l[acc]: + acc += 1 + corr = corr_l[acc] if acc < K else None + clean, new_specs = samp_l[0], samp_l[1:] + new_q = P[n:blk] + if corr is not None: # rejected at spec `acc` + rt.commit(cache, 1 + acc) + stats.accept_hist[acc] += 1 + out = emit(specs[:acc] + [corr]) + t0, t0_emitted = corr, True + pending, specs, spec_q = None, None, None + else: # all K accepted: N tokens this forward + rt.commit(cache, n) + stats.accept_hist[K] += 1 + out = emit(specs + [clean]) + pending, specs, spec_q = clean, new_specs, new_q + else: # ---- cold start + canvas = [t0] + [mask_id] * (blk - 1) + L = rt.canvas(cache, canvas, n, rows=n if rt.skip_rows else None) + stats.forwards += 1 + stats.cold_forwards += 1 + if sp.greedy: + vals = mx.argmax(L[:n], axis=-1).tolist() + new_q = None + else: + P = truncated_probs(L[:n], sp) + vals = sample_from_probs(P).tolist() + new_q = P[1:n] + clean, new_specs = vals[0], vals[1:n] + rt.commit(cache, 1) + out = emit(([] if t0_emitted else [t0]) + [clean]) + t0_emitted = True + pending, specs, spec_q = clean, new_specs, new_q + if out: + yield out + stats.decode_s = time.perf_counter() - t_start diff --git a/inference/mlx/trida_mlx/engine.py b/inference/mlx/trida_mlx/engine.py new file mode 100644 index 0000000..89e338c --- /dev/null +++ b/inference/mlx/trida_mlx/engine.py @@ -0,0 +1,228 @@ +"""Engine: model + tokenizer + chat template + a single-sequence prompt cache. + +The prompt cache is what makes a local agent loop fast: each turn re-sends the whole +conversation, but only the new messages need a prefill. Gated-delta layers cannot rewind +an arbitrary number of tokens, so the engine keeps two resumable points: + + * the live cache (everything committed at the end of the last generation), and + * a snapshot taken right before the last ``<|im_start|>`` of the previous prompt + (i.e. before the ``assistant`` generation header), which is where the next turn's + re-rendered history is guaranteed to still agree with it. +""" +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path +from typing import Iterator, Optional + +import mlx.core as mx + +from .decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate +from .model import DEFAULT_MODEL, TridaRuntime, load_model, mask_id_of + + +class IncrementalDetokenizer: + def __init__(self, tokenizer): + self.tok = tokenizer + self.ids: list = [] + self.prefix = 0 + self.read = 0 + + def add(self, new_ids: list) -> str: + self.ids.extend(new_ids) + prefix_text = self.tok.decode(self.ids[self.prefix : self.read], skip_special_tokens=False) + new_text = self.tok.decode(self.ids[self.prefix :], skip_special_tokens=False) + if new_text.endswith("�"): + return "" + delta = new_text[len(prefix_text):] + self.prefix, self.read = self.read, len(self.ids) + return delta + + +class _Slot: + """One conversation's cache plus resumable snapshots: ``snaps['system']`` right after the + system block (stable across every turn of an agent loop, and often 10k+ tokens of tool + schemas) and ``snaps['last']`` before the last ``<|im_start|>`` of the previous prompt.""" + __slots__ = ("cache", "snaps", "used") + + def __init__(self, cache): + self.cache, self.snaps, self.used = cache, {}, time.monotonic() + + +class Engine: + """``cache_slots``: how many independent conversations keep a resumable cache. An agent + harness interleaves its main loop with side requests (titles, summaries, memory review); + with one slot every side request would evict the main conversation's cache.""" + + def __init__(self, model: str = DEFAULT_MODEL, *, mode: str = "self-spec", gen_block: int = 4, + prefill_chunk: int = 512, prompt_cache: bool = True, fused_gdn: bool = True, + cache_slots: int = 3): + from transformers import AutoTokenizer + + t = time.perf_counter() + self.model, self.cfg, self.path = load_model(model) + self.tokenizer = AutoTokenizer.from_pretrained(str(self.path)) + self.mask_id = mask_id_of(self.path, self.cfg, self.tokenizer) + self.rt = TridaRuntime(self.model, prefill_chunk=prefill_chunk, fused_gdn=fused_gdn) + self.mode, self.n = mode, gen_block + self.model_id = model + self.eos_ids = self._eos_ids() + self.im_start = self.tokenizer.convert_tokens_to_ids("<|im_start|>") + self.use_prompt_cache = prompt_cache + self.cache_slots = max(1, cache_slots) + self.reset_cache() + self.lock = threading.Lock() + self.last_stats: Optional[dict] = None + self.load_s = time.perf_counter() - t + + # current slot's cache / snapshot + @property + def cache(self): + return self.cur.cache + + @cache.setter + def cache(self, c): + self.cur.cache = c + + def reset_cache(self) -> None: + self.cur = _Slot(self.rt.make_cache()) + self.slots = [self.cur] + + def cached_tokens(self) -> list: + return [s.cache.length for s in self.slots] + def _eos_ids(self) -> set: + ids = set() + eos = self.tokenizer.eos_token_id + if eos is not None: + ids.add(int(eos)) + gc = self.path / "generation_config.json" + if gc.is_file(): + e = json.loads(gc.read_text()).get("eos_token_id") + ids.update([e] if isinstance(e, int) else (e or [])) + for t in ("<|im_end|>", "<|endoftext|>"): + tid = self.tokenizer.convert_tokens_to_ids(t) + if isinstance(tid, int) and tid >= 0 and tid != self.tokenizer.unk_token_id: + ids.add(tid) + return ids + + # -- prompts ---------------------------------------------------------------- + def render(self, messages: list, tools: Optional[list] = None, enable_thinking: bool = True, + **template_kwargs) -> str: + return self.tokenizer.apply_chat_template( + messages, tools=tools or None, add_generation_prompt=True, tokenize=False, + enable_thinking=enable_thinking, **template_kwargs, + ) + + def encode(self, text: str) -> list: + return self.tokenizer.encode(text, add_special_tokens=False) + + # -- prompt cache ----------------------------------------------------------- + def _position_cache(self, prompt: list) -> None: + """Point ``self.cur`` at the slot holding the longest reusable prefix of ``prompt`` + (always < len(prompt)); otherwise take a fresh slot (evicting the least recently used).""" + def usable(tokens): + return 0 < len(tokens) < len(prompt) and prompt[: len(tokens)] == tokens + + if not self.use_prompt_cache: + self.reset_cache() + return + best, best_slot, best_kind = 0, None, None + for sl in self.slots: + live = sl.cache.tokens + if usable(live) and len(live) > best: + best, best_slot, best_kind = len(live), sl, None + for snap in sl.snaps.values(): + if usable(snap[0]) and len(snap[0]) > best: + best, best_slot, best_kind = len(snap[0]), sl, snap + if best_slot is not None: + self.cur = best_slot + if best_kind is not None: + best_slot.cache.restore(best_kind) + # drop snapshots that are no longer a prefix of what the cache now holds + held = best_slot.cache.tokens + best_slot.snaps = {k: v for k, v in best_slot.snaps.items() + if len(v[0]) <= len(held) and held[: len(v[0])] == v[0]} + elif any(sl.cache.length == 0 for sl in self.slots): + self.cur = next(sl for sl in self.slots if sl.cache.length == 0) + elif len(self.slots) < self.cache_slots: + self.cur = _Slot(self.rt.make_cache()) + self.slots.append(self.cur) + else: + self.cur = min(self.slots, key=lambda s: s.used) + self.cur.cache, self.cur.snaps = self.rt.make_cache(), {} + self.cur.used = time.monotonic() + + def _prefill_to_snapshot(self, prompt: list) -> None: + if not self.use_prompt_cache: + return + starts = [i for i, t in enumerate(prompt) if t == self.im_start] + marks = [] + if len(starts) >= 2 and starts[0] == 0: + marks.append(("system", starts[1])) # end of the system block + if starts: + marks.append(("last", starts[-1])) # before the generation header + for name, cut in marks: + if cut > self.cache.length: + self.rt.prefill(self.cache, prompt[self.cache.length : cut]) + self.cur.snaps[name] = self.cache.snapshot() + + # -- generation ------------------------------------------------------------- + def generate_ids(self, prompt: list, *, max_tokens: int = 2048, sampling: Optional[SamplingParams] = None, + mode: Optional[str] = None, stats: Optional[DecodeStats] = None, + progress=None) -> Iterator[list]: + """Yields lists of new token ids. Caller must hold ``self.lock``. + ``progress(done, total)`` is called between prefill chunks (keep-alives for clients).""" + sp = sampling or SamplingParams() + mode = mode or self.mode + stats = stats if stats is not None else DecodeStats() + t = time.perf_counter() + self._position_cache(prompt) + reused = self.cache.length + if progress is not None: + done = [0] + total = len(prompt) - reused + + def _cb(k): + done[0] += k + progress(done[0], total) + self.rt.progress_cb = _cb + try: + self._prefill_to_snapshot(prompt) + except BaseException: + self.rt.progress_cb = None + raise + pre_s = time.perf_counter() - t + if mode in ("causal", "ar"): + gen = ar_generate(self.rt, self.cache, prompt, max_new_tokens=max_tokens, eos_ids=self.eos_ids, + sp=sp, stats=stats) + else: + gen = selfspec_generate(self.rt, self.cache, prompt, max_new_tokens=max_tokens, + eos_ids=self.eos_ids, sp=sp, mask_id=self.mask_id, n=self.n, stats=stats) + try: + for i, chunk in enumerate(gen): + if i == 0: + self.rt.progress_cb = None + yield chunk + finally: + gen.close() + self.rt.progress_cb = None + stats.prefill_s += pre_s + stats.reused_tokens = reused + self.last_stats = stats.as_dict() + + def generate_text(self, prompt_text: str, **kw) -> Iterator[str]: + detok = IncrementalDetokenizer(self.tokenizer) + for ids in self.generate_ids(self.encode(prompt_text), **kw): + d = detok.add(ids) + if d: + yield d + + def warmup(self) -> None: + with self.lock: + p = self.encode("<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\n") + for _ in self.generate_ids(p, max_tokens=8, sampling=SamplingParams(temperature=0.0)): + pass + self.reset_cache() + mx.clear_cache() diff --git a/inference/mlx/trida_mlx/kernels.py b/inference/mlx/trida_mlx/kernels.py new file mode 100644 index 0000000..54a7447 --- /dev/null +++ b/inference/mlx/trida_mlx/kernels.py @@ -0,0 +1,133 @@ +"""Fused Metal kernel for the self-spec canvas gated-delta step. + +One launch per GDN layer replaces (kernel over clean rows) + (kernel over MASK rows) + +(repeat + einsum block-end readout): + + for t in 0..T-1: S = g_t * S ; S += k_t (v_t - S k_t) beta_t (the recurrence) + t < NC : y_t = S q_t (token-causal clean rows) + t == NC-1: s_clean = S (state after the clean rows, for commit) + for t in NC..T-1: y_t = S_end q_t (MASK rows read the block-end state) + +The per-step arithmetic of the clean rows is the same as mlx-lm's gated_delta kernel (same +thread layout, same float accumulation order), so the clean rows stay bit-compatible with +the AR step. Derived from mlx-lm's ``gated_delta.py`` kernel (MIT, Apple Inc.). +""" +from __future__ import annotations + +import mlx.core as mx + +_SOURCE = """ + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto y_ = y + b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_clean = s_clean + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + auto g_ = g + b_idx * T * Hv; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + float kv_mem = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] * g_[hv_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (t < NC && thread_index_in_simdgroup == 0) { + y_[dv_idx] = static_cast(out); + } + if (t == NC - 1) { + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_clean[s_idx] = static_cast(state[i]); + } + } + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y_ += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + + // MASK rows: read out the block-end state + auto q2 = q + b_idx * T * Hk * Dk + hk_idx * Dk + NC * Hk * Dk; + auto y2 = y + b_idx * T * Hv * Dv + hv_idx * Dv + NC * Hv * Dv; + for (int t = NC; t < T; ++t) { + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + out += state[i] * q2[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y2[dv_idx] = static_cast(out); + } + q2 += Hk * Dk; + y2 += Hv * Dv; + } +""" + +_KERNEL = None + + +def _kernel(): + global _KERNEL + if _KERNEL is None: + _KERNEL = mx.fast.metal_kernel( + name="trida_selfspec_canvas_gdn", + input_names=["q", "k", "v", "g", "beta", "state_in", "T"], + output_names=["y", "s_clean"], + source=_SOURCE, + ) + return _KERNEL + + +def fused_available() -> bool: + return mx.default_device() == mx.gpu and mx.metal.is_available() + + +def canvas_gdn_fused(q, k, v, g, beta, state, n_clean: int): + """q,k: [B,T,Hk,Dk]; v: [B,T,Hv,Dv]; g,beta: [B,T,Hv]; state: [B,Hv,Dv,Dk] (float32). + Returns (y [B,T,Hv,Dv], s_clean [B,Hv,Dv,Dk]).""" + B, T, Hk, Dk = k.shape + Hv, Dv = v.shape[2:] + if Dk % 32: + raise ValueError("fused canvas GDN kernel needs Dk % 32 == 0") + return _kernel()( + inputs=[q, k, v, g, beta, state, T], + template=[("InT", q.dtype), ("StT", state.dtype), ("Dk", Dk), ("Dv", Dv), ("Hk", Hk), ("Hv", Hv), + ("NC", n_clean)], + grid=(32, Dv, B * Hv), + threadgroup=(32, 4, 1), + output_shapes=[(B, T, Hv, Dv), state.shape], + output_dtypes=[q.dtype, state.dtype], + ) diff --git a/inference/mlx/trida_mlx/model.py b/inference/mlx/trida_mlx/model.py new file mode 100644 index 0000000..17249ad --- /dev/null +++ b/inference/mlx/trida_mlx/model.py @@ -0,0 +1,403 @@ +"""Trida2.0-4B (Qwen3.5 hybrid: gated-delta linear attention + full attention) on MLX. + +This module owns the *forward passes* the self-speculative decoder needs, built on the +weights/modules of ``mlx_lm.models.qwen3_5`` (so quantized checkpoints, the Metal +gated-delta kernel and RoPE all come from mlx-lm unchanged): + +* ``prefill`` causal prompt prefill (chunked), persists KV + GDN state. +* ``ar_step`` one causal token (the AR baseline, ``--mode causal``). +* ``canvas`` one self-spec forward over a ``2N-1`` token canvas with the + ``bd_bidir_shift`` semantics of the SGLang reference + (``HybridDiffusionSelfSpec``): + + - full attention: rows ``0..N-1`` are token-causal, rows ``N..2N-2`` (the MASK + slots) attend to the whole canvas; the prefix is visible to every row. + - gated-delta (``causal_mode=2, num_clean=N``): rows ``0..N-1`` read the recurrent + state *at their own position*; MASK rows read the state at the *block end*. + - nothing is persisted by the forward: ``commit(adv)`` later keeps the first + ``adv`` canvas tokens (KV trim + GDN state/conv window of the accepted prefix). + +Because rows ``0..N-1`` are exactly the causal computation, their logits are the AR +logits, which is what makes greedy self-spec lossless w.r.t. ``ar_step``. +""" +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models import qwen3_5 +from mlx_lm.models.cache import KVCache +from mlx_lm.models.gated_delta import compute_g, gated_delta_kernel, gated_delta_ops + +from .kernels import canvas_gdn_fused, fused_available + +DEFAULT_MODEL = "trillionlabs/Trida2.0-4B" +DEFAULT_MASK_ID = 248077 + + +# ----------------------------------------------------------------------------- loading +def resolve_model_path(path_or_id: str) -> Path: + p = Path(path_or_id).expanduser() + if p.is_dir(): + return p + if path_or_id.startswith(("/", "~", ".")) or path_or_id.count("/") > 1: + raise FileNotFoundError(f"local model directory not found: {p}") + from huggingface_hub import snapshot_download + + return Path( + snapshot_download( + path_or_id, + allow_patterns=["*.json", "*.safetensors", "*.jinja", "*.txt", "*.model", "*.tiktoken"], + ) + ) + + +def _text_config(cfg: dict) -> dict: + text = dict(cfg.get("text_config", cfg)) + text.pop("architectures", None) + return text + + +def build_model(text_cfg: dict) -> qwen3_5.Model: + args = qwen3_5.ModelArgs(model_type="qwen3_5", text_config=dict(text_cfg)) + return qwen3_5.Model(args) + + +def load_model(path_or_id: str = DEFAULT_MODEL, *, dtype=mx.bfloat16): + """Load a raw HF Trida checkpoint or one written by ``convert.py`` (optionally quantized). + + Returns ``(model, config_dict, path)``. + """ + path = resolve_model_path(path_or_id) + cfg = json.loads((path / "config.json").read_text()) + model = build_model(_text_config(cfg)) + weights = {} + for f in sorted(path.glob("*.safetensors")): + weights.update(mx.load(str(f))) + if not weights: + raise FileNotFoundError(f"no *.safetensors in {path}") + quant = cfg.get("quantization") + if quant: + nn.quantize( + model, + group_size=quant["group_size"], + bits=quant["bits"], + class_predicate=lambda p, m: f"{p}.scales" in weights, + ) + weights = model.sanitize(weights) + if model.language_model.args.tie_word_embeddings: + # mlx-lm's qwen3_5 sanitize pops "lm_head.weight" before it is prefixed, so a raw + # text-only checkpoint that also ships the (tied) head leaves it behind. With tied + # embeddings HF/vLLM ignore it and use embed_tokens; do the same. + head = weights.pop("language_model.lm_head.weight", None) + for sfx in ("scales", "biases"): + weights.pop(f"language_model.lm_head.{sfx}", None) + emb = weights.get("language_model.model.embed_tokens.weight") + if head is not None and emb is not None and head.shape == emb.shape and not quant: + if not mx.array_equal(head.astype(mx.float32), emb.astype(mx.float32)).item(): + print("[trida-mlx] warning: config ties embeddings but lm_head.weight differs from " + "embed_tokens; using embed_tokens (same as HF/vLLM)") + if not quant: + weights = { + k: (v if k.endswith("A_log") or not mx.issubdtype(v.dtype, mx.floating) else v.astype(dtype)) + for k, v in weights.items() + } + model.load_weights(list(weights.items()), strict=True) + mx.eval(model.parameters()) + model.eval() + return model, cfg, path + + +def mask_id_of(path: Path, cfg: dict, tokenizer=None) -> int: + bd = path / "block_diffusion.json" + if bd.is_file(): + return int(json.loads(bd.read_text())["mask_id"]) + if "mask_id" in cfg: + return int(cfg["mask_id"]) + if tokenizer is not None: + tid = tokenizer.convert_tokens_to_ids("<|mask|>") + if isinstance(tid, int) and tid >= 0 and tid != tokenizer.unk_token_id: + return tid + return DEFAULT_MASK_ID + + +# ----------------------------------------------------------------------------- cache +@dataclass +class GDNState: + conv: mx.array # [1, K-1, conv_dim] last pre-conv rows + ssm: mx.array # [1, Hv, Dv, Dk] float32 recurrent state + + +@dataclass +class _CanvasGDN: + """What a canvas forward leaves behind per GDN layer so ``commit`` can persist + the state after any accepted prefix without re-running the model.""" + + prev: GDNState + conv_input: mx.array # [1, K-1+T, conv_dim] = prev conv rows ++ canvas pre-conv rows + q: mx.array + k: mx.array + v: mx.array + g: mx.array + beta: mx.array # clean rows only (first n_clean) + s_clean: mx.array # state after all n_clean clean rows + + +@dataclass +class SeqCache: + """Per-sequence decode state: KV caches (attention layers), GDN states (linear + layers) and the list of committed token ids (what the state encodes).""" + + layers: list + tokens: list = field(default_factory=list) + pending: Optional[list] = None # _CanvasGDN per layer after a canvas forward + canvas_len: int = 0 + + @property + def length(self) -> int: + return len(self.tokens) + + # A snapshot is cheap: MLX arrays are immutable, so GDN states are kept by + # reference; attention KV only needs its length (it is append-only). + def snapshot(self): + assert self.pending is None + return ( + list(self.tokens), + [l if isinstance(l, GDNState) else l.offset for l in self.layers], + ) + + def restore(self, snap) -> None: + tokens, states = snap + assert self.pending is None + for i, s in enumerate(states): + if isinstance(s, GDNState): + self.layers[i] = s + else: + kv = self.layers[i] + if s > kv.offset: + raise ValueError("cannot restore a KV snapshot longer than the cache") + kv.trim(kv.offset - s) + self.tokens = list(tokens) + + def nbytes(self) -> int: + n = 0 + for l in self.layers: + if isinstance(l, GDNState): + n += l.conv.nbytes + l.ssm.nbytes + else: + n += l.nbytes + return n + + +# ----------------------------------------------------------------------------- runtime +def _gdr(q, k, v, g, beta, state): + """Gated-delta recurrence: Metal kernel on GPU, reference ops elsewhere.""" + if mx.default_device() == mx.gpu and mx.metal.is_available(): + return gated_delta_kernel(q, k, v, g, beta, state) + return gated_delta_ops(q, k, v, g, beta, state) + + +class TridaRuntime: + def __init__(self, model: qwen3_5.Model, prefill_chunk: int = 512, fused_gdn: bool = True): + self.model = model + # one Metal launch per GDN layer for the canvas (GPU only; ops fallback elsewhere) + self.fused_gdn = fused_gdn + # cold start: LM head on rows 0..N-1 only. Off by default: measured slower on M3 Pro + # (46.2 -> 48.3 ms) since the head is bandwidth-bound and rows are nearly free. + self.skip_rows = False + self.lm = model.language_model # TextModel + self.backbone = self.lm.model # Qwen3_5TextModel + self.args = self.lm.args + self.prefill_chunk = prefill_chunk + self.progress_cb = None # optional fn(tokens_done_this_call) during long prefills + self._mask_cache: dict = {} + + # -- cache ------------------------------------------------------------------ + def make_cache(self) -> SeqCache: + layers = [] + for layer in self.backbone.layers: + if layer.is_linear: + lin = layer.linear_attn + dtype = self.backbone.norm.weight.dtype + layers.append( + GDNState( + conv=mx.zeros((1, lin.conv_kernel_size - 1, lin.conv_dim), dtype=dtype), + ssm=mx.zeros((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim), dtype=mx.float32), + ) + ) + else: + layers.append(KVCache()) + return SeqCache(layers=layers) + + # -- layers ----------------------------------------------------------------- + @staticmethod + def _attention(attn, x, kv: KVCache, mask): + B, L, _ = x.shape + qo = attn.q_proj(x) + queries, gate = mx.split(qo.reshape(B, L, attn.num_attention_heads, -1), 2, axis=-1) + gate = gate.reshape(B, L, -1) + keys, values = attn.k_proj(x), attn.v_proj(x) + queries = attn.q_norm(queries).transpose(0, 2, 1, 3) + keys = attn.k_norm(keys.reshape(B, L, attn.num_key_value_heads, -1)).transpose(0, 2, 1, 3) + values = values.reshape(B, L, attn.num_key_value_heads, -1).transpose(0, 2, 1, 3) + queries = attn.rope(queries, offset=kv.offset) + keys = attn.rope(keys, offset=kv.offset) + keys, values = kv.update_and_fetch(keys, values) + out = mx.fast.scaled_dot_product_attention(queries, keys, values, scale=attn.scale, mask=mask) + out = out.transpose(0, 2, 1, 3).reshape(B, L, -1) + return attn.o_proj(out * mx.sigmoid(gate)) + + def _gdn(self, lin, x, st: GDNState, n_clean: Optional[int]): + """Returns (out, new_state_or_None, canvas_record_or_None).""" + B, S, _ = x.shape + qkv = lin.in_proj_qkv(x) + z = lin.in_proj_z(x).reshape(B, S, lin.num_v_heads, lin.head_v_dim) + b = lin.in_proj_b(x) + a = lin.in_proj_a(x) + conv_input = mx.concatenate([st.conv.astype(qkv.dtype), qkv], axis=1) + conv_out = nn.silu(lin.conv1d(conv_input)) + q, k, v = [ + t.reshape(B, S, h, d) + for t, h, d in zip( + mx.split(conv_out, [lin.key_dim, 2 * lin.key_dim], -1), + [lin.num_k_heads, lin.num_k_heads, lin.num_v_heads], + [lin.head_k_dim, lin.head_k_dim, lin.head_v_dim], + ) + ] + inv_scale = k.shape[-1] ** -0.5 + q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6) + k = inv_scale * mx.fast.rms_norm(k, None, 1e-6) + beta = mx.sigmoid(b) + g = compute_g(lin.A_log, a, lin.dt_bias) + + if n_clean is None: # plain causal (prefill / AR step) + y, s = _gdr(q, k, v, g, beta, st.ssm) + n_keep = lin.conv_kernel_size - 1 + new = GDNState(conv=conv_input[:, -n_keep:], ssm=s) + rec = None + else: # canvas: clean rows token-causal, MASK rows read the block-end state + n = n_clean + if self.fused_gdn and fused_available(): + y, s_clean = canvas_gdn_fused(q, k, v, g, beta, st.ssm, n) + else: + y1, s_clean = _gdr(q[:, :n], k[:, :n], v[:, :n], g[:, :n], beta[:, :n], st.ssm) + _, s_end = _gdr(q[:, n:], k[:, n:], v[:, n:], g[:, n:], beta[:, n:], s_clean) + qm = q[:, n:] + rep = lin.num_v_heads // lin.num_k_heads + if rep > 1: + qm = mx.repeat(qm, rep, axis=2) + y2 = mx.einsum("bhvk,bthk->bthv", s_end, qm.astype(mx.float32)).astype(y1.dtype) + y = mx.concatenate([y1, y2], axis=1) + new = None + rec = _CanvasGDN( + prev=st, conv_input=conv_input, + q=q[:, :n], k=k[:, :n], v=v[:, :n], g=g[:, :n], beta=beta[:, :n], + s_clean=s_clean, + ) + out = lin.norm(y, z) + return lin.out_proj(out.reshape(B, S, -1)), new, rec + + def _lm_head(self, h): + if self.args.tie_word_embeddings: + return self.backbone.embed_tokens.as_linear(h) + return self.lm.lm_head(h) + + def _forward(self, ids: mx.array, cache: SeqCache, attn_mask, n_clean, last_only: bool, + rows: Optional[int] = None): + h = self.backbone.embed_tokens(ids) + recs = [] if n_clean is not None else None + for i, layer in enumerate(self.backbone.layers): + x = layer.input_layernorm(h) + c = cache.layers[i] + if layer.is_linear: + r, new, rec = self._gdn(layer.linear_attn, x, c, n_clean) + if new is not None: + cache.layers[i] = new + if recs is not None: + recs.append(rec) + else: + r = self._attention(layer.self_attn, x, c, attn_mask) + if recs is not None: + recs.append(None) + h = h + r + h = h + layer.mlp(layer.post_attention_layernorm(h)) + h = self.backbone.norm(h) + if last_only: + h = h[:, -1:] + elif rows is not None: + h = h[:, :rows] + return self._lm_head(h), recs + + # -- public forwards -------------------------------------------------------- + def prefill(self, cache: SeqCache, tokens: list) -> mx.array: + """Causal prefill of ``tokens`` after whatever ``cache`` already holds. + Returns the float32 logits [V] of the last token.""" + assert cache.pending is None and tokens + logits = None + for s in range(0, len(tokens), self.prefill_chunk): + chunk = tokens[s : s + self.prefill_chunk] + ids = mx.array([chunk], dtype=mx.int32) + mask = "causal" if len(chunk) > 1 else None + logits, _ = self._forward(ids, cache, mask, None, last_only=True) + cache.tokens.extend(chunk) + if s + self.prefill_chunk < len(tokens): + mx.eval([l.ssm if isinstance(l, GDNState) else l.keys for l in cache.layers]) + if self.progress_cb is not None: + self.progress_cb(len(chunk)) + return logits[0, -1].astype(mx.float32) + + def ar_step(self, cache: SeqCache, token: int) -> mx.array: + logits, _ = self._forward(mx.array([[token]], dtype=mx.int32), cache, None, None, last_only=True) + cache.tokens.append(token) + return logits[0, -1].astype(mx.float32) + + def _canvas_mask(self, prefix_len: int, blk: int, n_clean: int) -> mx.array: + key = (blk, n_clean) + inner = self._mask_cache.get(key) + if inner is None: + r = mx.arange(blk)[:, None] + c = mx.arange(blk)[None, :] + inner = (c <= r) | (r >= n_clean) # clean rows causal, MASK rows see the whole canvas + self._mask_cache[key] = inner + pre = mx.ones((blk, prefix_len), dtype=mx.bool_) + return mx.concatenate([pre, inner], axis=1) + + def canvas(self, cache: SeqCache, canvas_tokens: list, n_clean: int, rows: Optional[int] = None) -> mx.array: + """One self-spec forward. Returns float32 logits [rows or blk, V] (shift readout: row i + predicts the token at canvas position i+1). ``rows`` limits the LM head to the first + rows (a cold start only reads rows 0..N-1). Call ``commit`` afterwards.""" + assert cache.pending is None + blk = len(canvas_tokens) + mask = self._canvas_mask(cache.length, blk, n_clean) + ids = mx.array([canvas_tokens], dtype=mx.int32) + logits, recs = self._forward(ids, cache, mask, n_clean, last_only=False, rows=rows) + cache.pending = recs + cache.canvas_len = blk + self._canvas_tokens = list(canvas_tokens) + return logits[0].astype(mx.float32) + + def commit(self, cache: SeqCache, adv: int) -> None: + """Keep the first ``adv`` (1..n_clean) canvas tokens in the cache.""" + recs, blk = cache.pending, cache.canvas_len + assert recs is not None and 1 <= adv + for i, rec in enumerate(recs): + if rec is None: + cache.layers[i].trim(blk - adv) + continue + n = rec.q.shape[1] + assert adv <= n + n_keep = rec.conv_input.shape[1] - blk + conv = rec.conv_input[:, adv : adv + n_keep] + if adv == n: + ssm = rec.s_clean + else: + _, ssm = _gdr(rec.q[:, :adv], rec.k[:, :adv], rec.v[:, :adv], + rec.g[:, :adv], rec.beta[:, :adv], rec.prev.ssm) + cache.layers[i] = GDNState(conv=conv, ssm=ssm) + cache.tokens.extend(self._canvas_tokens[:adv]) + cache.pending = None + cache.canvas_len = 0 diff --git a/inference/mlx/trida_mlx/profile.py b/inference/mlx/trida_mlx/profile.py new file mode 100644 index 0000000..ec4518f --- /dev/null +++ b/inference/mlx/trida_mlx/profile.py @@ -0,0 +1,166 @@ +"""Where does a self-spec step spend its time? (run on the Mac) + + uv run trida-mlx-profile --model ~/models/Trida2.0-4B-mlx-q8 + +Times, at the same KV length: the AR step (1 token) vs the self-spec canvas step +(2N-1 tokens), each split into LM head / MLPs / full attention / gated-delta, plus the +Python graph-build time and the commit, and A/Bs the step optimizations (fused canvas GDN +kernel, LM-head row skipping on cold starts). Medians over --iters runs, GPU-synchronized. +""" +from __future__ import annotations + +import argparse +import statistics +import time + +import mlx.core as mx + +from .engine import Engine +from .model import GDNState + +LONG_PROMPT = [{"role": "user", "content": "Summarize the history of computing in detail. " * 40}] + + +def timeit(fn, iters, warmup=3): + for _ in range(warmup): + mx.eval(fn()) + ts = [] + for _ in range(iters): + t = time.perf_counter() + mx.eval(fn()) + ts.append(time.perf_counter() - t) + return statistics.median(ts) * 1e3 + + +def main(argv=None): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") + ap.add_argument("--gen-block", type=int, default=4) + ap.add_argument("--iters", type=int, default=20) + a = ap.parse_args(argv) + eng = Engine(a.model, gen_block=a.gen_block, prompt_cache=False) + rt, n, mask = eng.rt, a.gen_block, eng.mask_id + blk = 2 * n - 1 + prompt = eng.encode(eng.render(LONG_PROMPT, enable_thinking=False)) + cache = rt.make_cache() + rt.prefill(cache, prompt) + mx.eval([l.ssm if isinstance(l, GDNState) else l.keys for l in cache.layers]) + P = cache.length + D = rt.args.hidden_size + layers = rt.backbone.layers + print(f"prefix length {P}, canvas {blk} tokens, device {mx.default_device()}") + + rows = [] + + def row(name, ms1, ms7, note=""): + rows.append((name, ms1, ms7, note)) + + # ---- whole steps ------------------------------------------------------- + def ar_step(): + logits = rt.ar_step(cache, 42) + tok = mx.argmax(logits) + mx.eval(tok) + return tok + + def trim_ar(): + # undo the appended token so every run sees the same prefix + for i, l in enumerate(cache.layers): + if not isinstance(l, GDNState): + l.trim(l.offset - P) + cache.tokens = cache.tokens[:P] + + snap_states = [l for l in cache.layers] + + def restore(): + trim_ar() + for i, s in enumerate(snap_states): + if isinstance(s, GDNState): + cache.layers[i] = s + + def run_ar(): + restore() + return ar_step() + + def run_canvas(rows_=None, adv=n): + restore() + L = rt.canvas(cache, [11, 12, 13, 14][:n] + [mask] * (n - 1), n, rows=rows_) + am = mx.argmax(L, axis=-1) + mx.eval(am) + rt.commit(cache, adv) + return [l.ssm if isinstance(l, GDNState) else l.keys for l in cache.layers] + + t_ar = timeit(run_ar, a.iters) + rt.fused_gdn = False + t_cv_unfused = timeit(run_canvas, a.iters) + rt.fused_gdn = True + t_cv = timeit(run_canvas, a.iters) + t_cold = timeit(lambda: run_canvas(rows_=n, adv=1), a.iters) + t_cold_all = timeit(lambda: run_canvas(rows_=None, adv=1), a.iters) + t_rej = timeit(lambda: run_canvas(adv=2), a.iters) + + # graph-build (host) time only + def build_only(): + restore() + t = time.perf_counter() + L = rt.canvas(cache, [11, 12, 13, 14][:n] + [mask] * (n - 1), n) + dt = time.perf_counter() - t + mx.eval(L) + cache.pending = None + return dt + for _ in range(3): + build_only() + t_build = statistics.median(build_only() for _ in range(a.iters)) * 1e3 + restore() + + # ---- components ---------------------------------------------------------- + def comp(L): + x = mx.random.normal((1, L, D)).astype(mx.bfloat16) + mx.eval(x) + t_head = timeit(lambda: rt._lm_head(x), a.iters) + t_mlp = timeit(lambda: [l.mlp(l.post_attention_layernorm(x)) for l in layers], a.iters) + + def attn(): + outs = [] + for i, l in enumerate(layers): + if not l.is_linear: + kv = cache.layers[i] + m = rt._canvas_mask(kv.offset, L, n) if L > 1 else None + outs.append(rt._attention(l.self_attn, x, kv, m)) + kv.trim(L) + return outs + t_attn = timeit(attn, a.iters) + + def gdn(): + return [rt._gdn(l.linear_attn, x, cache.layers[i], n if L > 1 else None)[0] + for i, l in enumerate(layers) if l.is_linear] + t_gdn = timeit(gdn, a.iters) + return t_head, t_mlp, t_attn, t_gdn + + c1 = comp(1) + c7 = comp(blk) + rt.fused_gdn = False + c7u = comp(blk) + rt.fused_gdn = True + + print(f"\n{'':34s} {'AR (1 tok)':>11s} {'canvas (%d)' % blk:>11s} ratio") + print(f"{'whole step (fwd+argmax+commit)':34s} {t_ar:9.2f}ms {t_cv:9.2f}ms {t_cv / t_ar:4.2f}x") + n_lin = sum(l.is_linear for l in layers) + for name, x1, x7 in [(" LM head", c1[0], c7[0]), (f" MLPs ({len(layers)})", c1[1], c7[1]), + (f" full attention ({len(layers) - n_lin})", c1[2], c7[2]), + (f" gated-delta ({n_lin})", c1[3], c7[3])]: + print(f"{name:34s} {x1:9.2f}ms {x7:9.2f}ms {x7 / x1:4.2f}x") + s1, s7 = sum(c1), sum(c7) + print(f"{' sum of components':34s} {s1:9.2f}ms {s7:9.2f}ms (rest = embed, norms, sync, commit)") + print(f"\ncanvas variants:") + print(f" verify step, fused GDN {t_cv:7.2f} ms") + print(f" verify step, unfused GDN {t_cv_unfused:7.2f} ms (GDN component {c7u[3]:.2f} ms vs fused {c7[3]:.2f} ms)") + print(f" cold step, LM head rows 0..{n - 1} {t_cold:7.2f} ms") + print(f" cold step, LM head all rows {t_cold_all:7.2f} ms") + print(f" verify + reject (commit adv=2) {t_rej:7.2f} ms (commit recompute cost {t_rej - t_cv:+.2f} ms)") + print(f" host graph build (canvas) {t_build:7.2f} ms") + print(f"\nbreak-even: self-spec beats AR when tokens/forward > {t_cv / t_ar:.2f}") + print(f"peak memory {mx.get_peak_memory() / 1e9:.2f} GB") + + +if __name__ == "__main__": + main() diff --git a/inference/mlx/trida_mlx/verify.py b/inference/mlx/trida_mlx/verify.py new file mode 100644 index 0000000..bb45629 --- /dev/null +++ b/inference/mlx/trida_mlx/verify.py @@ -0,0 +1,92 @@ +"""First-contact check for a real checkpoint on the Mac (run this before benchmarking). + + python -m trida_mlx.verify --model ./Trida2.0-4B-mlx-q8 + +1. prints the resolved config (layers, GDN/attention split, mask id, eos ids); +2. canvas vs AR logits: rows 0..N-1 of a verify canvas must equal step-by-step AR logits + (max |diff| should be at bf16 noise level, argmax identical); +3. greedy self-spec vs greedy AR for --max-tokens tokens. In bf16 the canvas (7-row matmuls) + and the AR step (1-row matmuls) use different kernels, so an exact tie-break can flip at a + near-tie; if outputs diverge, the AR top-1/top-2 logit margin at that position is printed + (a tiny margin = numerics, a large margin = a real bug). +""" +from __future__ import annotations + +import argparse +import time + +import mlx.core as mx + +from .decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate +from .engine import Engine + +PROMPT = [{"role": "user", "content": "Explain in a short paragraph why the sky is blue, then give a one-line summary."}] + + +def main(argv=None): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") + ap.add_argument("--max-tokens", type=int, default=256) + ap.add_argument("--gen-block", type=int, default=4) + a = ap.parse_args(argv) + eng = Engine(a.model, gen_block=a.gen_block, prompt_cache=False) + rt, n = eng.rt, a.gen_block + args = rt.args + n_lin = sum(l.is_linear for l in rt.backbone.layers) + print(f"model {a.model}: {args.num_hidden_layers} layers ({n_lin} gated-delta, {args.num_hidden_layers - n_lin} " + f"full-attn), hidden {args.hidden_size}, vocab {args.vocab_size}, tied={args.tie_word_embeddings}, " + f"quant={eng.cfg.get('quantization')}") + print(f"mask_id={eng.mask_id} ({eng.tokenizer.convert_ids_to_tokens(eng.mask_id)!r}) eos={sorted(eng.eos_ids)} " + f"load {eng.load_s:.1f}s device={mx.default_device()}") + + prompt = eng.encode(eng.render(PROMPT, enable_thinking=False)) + sp = SamplingParams(temperature=0.0) + + # 1) AR reference with margins + c = rt.make_cache() + logits = rt.prefill(c, prompt) + ar, margins, hit_eos = [], [], False + t = time.perf_counter() + for _ in range(a.max_tokens): + top2 = mx.topk(logits, 2) + tok = int(mx.argmax(logits).item()) + margins.append(float((mx.max(top2) - mx.min(top2)).item())) + if tok in eng.eos_ids: + hit_eos = True + break + ar.append(tok) + logits = rt.ar_step(c, tok) + ar_s = time.perf_counter() - t + + # 2) canvas rows vs AR logits on the first N generated tokens + c1 = rt.make_cache(); rt.prefill(c1, prompt) + L = rt.canvas(c1, ar[:n] + [eng.mask_id] * (n - 1), n) + c2 = rt.make_cache(); rt.prefill(c2, prompt) + worst, same = 0.0, True + for i in range(n): + ref = rt.ar_step(c2, ar[i]) + worst = max(worst, float(mx.max(mx.abs(L[i] - ref)).item())) + same &= int(mx.argmax(L[i]).item()) == int(mx.argmax(ref).item()) + print(f"[canvas==AR] rows 0..{n - 1}: max|dlogit|={worst:.4f} argmax identical={same}") + + # 3) greedy self-spec vs AR + st = DecodeStats() + t = time.perf_counter() + ss = [x for ch in selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=len(ar) + (1 if hit_eos else 0), + eos_ids=eng.eos_ids, sp=sp, mask_id=eng.mask_id, n=n, stats=st) for x in ch] + ss_s = time.perf_counter() - t + k = next((i for i in range(min(len(ar), len(ss))) if ar[i] != ss[i]), None) + if k is None and len(ar) == len(ss): + print(f"[lossless] YES: {len(ar)} tokens identical") + else: + k = k if k is not None else min(len(ar), len(ss)) + print(f"[lossless] diverged at token {k}/{len(ar)}; AR top1-top2 margin there = " + f"{margins[k] if k < len(margins) else float('nan'):.4f} (median margin {sorted(margins)[len(margins) // 2]:.3f})") + print(f"[speed] AR {len(ar) / ar_s:.1f} tok/s | self-spec {len(ss) / ss_s:.1f} tok/s " + f"({st.tokens_per_forward:.2f} tok/fwd, accept hist {st.accept_hist}, cold {st.cold_forwards}) " + f"peak mem {mx.get_peak_memory() / 1e9:.2f} GB") + print("---\n" + eng.tokenizer.decode(ss)[:600]) + + +if __name__ == "__main__": + main() diff --git a/inference/mlx/uv.lock b/inference/mlx/uv.lock new file mode 100644 index 0000000..fdcb1ff --- /dev/null +++ b/inference/mlx/uv.lock @@ -0,0 +1,783 @@ +version = 1 +revision = 3 +requires-python = ">=3.10, <3.14" +resolution-markers = [ + "python_full_version >= '3.12' and platform_machine == 'arm64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and platform_machine == 'arm64' and sys_platform == 'darwin'", + "python_full_version < '3.11' and platform_machine == 'arm64' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'linux'", +] +supported-markers = [ + "platform_machine == 'arm64' and sys_platform == 'darwin'", + "sys_platform == 'linux'", +] + +[[package]] +name = "annotated-doc" +version = "0.0.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/8e/38aa427ed5402449e226975b649c5dc73ccadfefeb95e6aecb8f8ea4b6b6/annotated_doc-0.0.5.tar.gz", hash = "sha256:c7e58ce09192557605d8bbd92836d7e1d520ac9580096042c0bfd197efacf1bb", size = 10758, upload-time = "2026-07-28T13:50:58.129Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3e/30/e900b21425a860e195f32e37657aa1f7c7f2b1bfb26f03ca209b90933c06/annotated_doc-0.0.5-py3-none-any.whl", hash = "sha256:117bac03a25ede5df5440e855b32d556049ca169ead221505badf432fed4b101", size = 5302, upload-time = "2026-07-28T13:50:57.239Z" }, +] + +[[package]] +name = "anyio" +version = "4.15.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "idna" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a9/d2/f4d173e22df740bc37b1db102b386ba719b66e95b0f0d751f556b387e6d2/anyio-4.15.1.tar.gz", hash = "sha256:9f28306018cbd6d329e64a36d58256edff76dd996fe423bc957326e578b82a94", size = 276966, upload-time = "2026-09-05T10:42:39.44Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/b8/4bd346e22b28902df4d651910f5242c28d84e4a5c2435ca5c3f797ed7e2e/anyio-4.15.1-py3-none-any.whl", hash = "sha256:6152fdbbf9a77fdec97731721bebf7c4c44f7c29b424b0065826173efc7ed101", size = 132079, upload-time = "2026-09-05T10:42:37.923Z" }, +] + +[[package]] +name = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "click" +version = "8.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/0e/7fa0ef50764b67090eca4114772a2abf8b6148198475e54c660b97caeee6/click-8.5.0.tar.gz", hash = "sha256:ba0d2089de75ea0310e2dde03160e6ca10009947fb95a182f9b54021bb272e34", size = 382235, upload-time = "2026-08-26T13:33:14.56Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/50/6c0d534c5f134586a8e1ba4e330569e32f057e33372ae556463212fb4cd3/click-8.5.0-py3-none-any.whl", hash = "sha256:255bc9599cf7748b4b1a446ccc735421bd08a2ae529a8b88597d3de5664ee360", size = 125251, upload-time = "2026-08-26T13:33:12.928Z" }, +] + +[[package]] +name = "exceptiongroup" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, +] + +[[package]] +name = "filelock" +version = "4.0.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/2d/5cb7a5ac017031e96a63173f3d57a8a5162ce36f9ff461e62bc6da4d88ab/filelock-4.0.6.tar.gz", hash = "sha256:323fab3b2fb22d889b29fa83774f60029addddb4b6a1bcfa1e73066eabffb5f2", size = 248237, upload-time = "2026-09-28T22:16:16.257Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/3b/b74ae9db2f78f1354278731372ae89870626be7b86cbb46b8d707acd1a98/filelock-4.0.6-py3-none-any.whl", hash = "sha256:9b139fb93b2ac5807f7feaf63aa4546fbd74074a2d3f554c040f40a4694d7b9f", size = 109077, upload-time = "2026-09-28T22:16:14.731Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/77/cd/9be253869fc42e764de7f3dedd6969af7d44ff9c3375214a3442a6f3fc08/fsspec-2026.9.0.tar.gz", hash = "sha256:0f08147951c8cb31d844c3547d631053b127863b60be04cf06e121333ee0e2fe", size = 333545, upload-time = "2026-09-18T17:50:42.825Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/c0/a98505f18594f1bce828bb159cec0fcf9860562f1a2c85913409fc8f3d9e/fsspec-2026.9.0-py3-none-any.whl", hash = "sha256:8dd6e646e99ea382bd85f97a45e6b526a442d79423a7dc673f1e2756d05fcb5f", size = 221738, upload-time = "2026-09-18T17:50:41.341Z" }, +] + +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, +] + +[[package]] +name = "hf-xet" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/ab/522a2ab67f27971a9d48ca666d4fca85ef7d5282d142e31fd087e27b1bbe/hf_xet-1.6.0.tar.gz", hash = "sha256:2e58454a340b3556dfa4972d5451aff4fba8dd42a236600ba1a1d2b1514f0fef", size = 920527, upload-time = "2026-08-03T22:33:13.243Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/69/55b8dcf636142ae660fec1869fcac14c4da2e8412e14d6eee1523be77e9f/hf_xet-1.6.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0906082d9932ae0c0057fa194041c22b4e2cdb46b2592ef3b91f020d62a081a", size = 3876287, upload-time = "2026-08-03T22:33:02.251Z" }, + { url = "https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d62671bb130879cef0ee4c9ebe47a14af6c66ec53e6d84dc15936e5ffdfac82f", size = 4464663, upload-time = "2026-08-03T22:33:03.802Z" }, + { url = "https://files.pythonhosted.org/packages/9a/69/1f0cbc2fb22ae6082d094f743d1b8945a3f36f6089cb95f42b7ee348cda7/hf_xet-1.6.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0e6e21fa3cdfcdcd76748564bf593870a5e013f47d97cf10aed63aa222cff5b7", size = 4262538, upload-time = "2026-08-03T22:33:05.287Z" }, + { url = "https://files.pythonhosted.org/packages/d1/3a/4f4f2301ade26e404462d3336fa11f7958d914cabbabdd6e03c3c5d5658c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:4fc74352a17015bd0ee90038bc9efe38db894cde45f268b6712b04fce8cd0acb", size = 4460520, upload-time = "2026-08-03T22:33:06.81Z" }, + { url = "https://files.pythonhosted.org/packages/ab/5f/311725e2a905534dfee2dcb5b08414f249147f1f12252bfc2bd24caa075c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8fb4f71cba6129110c3374a33f919001ff130488fc23553698e34cc1c2a1198c", size = 4675937, upload-time = "2026-08-03T22:33:08.616Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + +[[package]] +name = "huggingface-hub" +version = "1.33.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' or sys_platform != 'linux'" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/25/2a/484d112c0d8fc5f665d7b65137ac9cdb2953c982391598c3597968a12ee7/huggingface_hub-1.33.0.tar.gz", hash = "sha256:367be21a201db9523eddf8aeac7048f2602c1b308691c97640d5e72ed188007e", size = 1043548, upload-time = "2026-09-24T09:49:29.971Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/16/963096d224b80909432dc16561a615fd33d2d13beef3ce4c63fa25e40867/huggingface_hub-1.33.0-py3-none-any.whl", hash = "sha256:04e434b06e100eddbce9a6e817d72693a7884b10a79bd67ab48080d5c07eb899", size = 846435, upload-time = "2026-09-24T09:49:28.059Z" }, +] + +[[package]] +name = "idna" +version = "3.20" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/08/8eea9d4b8302028f3abb2c0813953f7aec26d33b7a8960ed760e65ff29fa/idna-3.20.tar.gz", hash = "sha256:a7db850025b95ded1eae8a46181a1a6c56c92c96f0e2b005d9ff8dc0210cab44", size = 216463, upload-time = "2026-09-17T14:11:04.752Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/a2/bb081bab032533a855d44de1d56f8e8426114ff1ba5d1f07a438a0a654f8/idna-3.20-py3-none-any.whl", hash = "sha256:ab7ae7122974553370f0bdb919e1a960b2cd1bc1ef0276416d896db81c14582c", size = 69583, upload-time = "2026-09-17T14:11:03.168Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/1b/fbd8eed11021cabd9226c37342fa6ca4e8a98d8188a8d9b66740494960e4/markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419", size = 12057, upload-time = "2025-09-27T18:36:07.165Z" }, + { url = "https://files.pythonhosted.org/packages/40/01/e560d658dc0bb8ab762670ece35281dec7b6c1b33f5fbc09ebb57a185519/markupsafe-3.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ba88449deb3de88bd40044603fafffb7bc2b055d626a330323a9ed736661695", size = 22050, upload-time = "2025-09-27T18:36:08.005Z" }, + { url = "https://files.pythonhosted.org/packages/af/cd/ce6e848bbf2c32314c9b237839119c5a564a59725b53157c856e90937b7a/markupsafe-3.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f42d0984e947b8adf7dd6dde396e720934d12c506ce84eea8476409563607591", size = 20681, upload-time = "2025-09-27T18:36:08.881Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2a/b5c12c809f1c3045c4d580b035a743d12fcde53cf685dbc44660826308da/markupsafe-3.0.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c0c0b3ade1c0b13b936d7970b1d37a57acde9199dc2aecc4c336773e1d86049c", size = 20705, upload-time = "2025-09-27T18:36:10.131Z" }, + { url = "https://files.pythonhosted.org/packages/cf/e3/9427a68c82728d0a88c50f890d0fc072a1484de2f3ac1ad0bfc1a7214fd5/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:0303439a41979d9e74d18ff5e2dd8c43ed6c6001fd40e5bf2e43f7bd9bbc523f", size = 21524, upload-time = "2025-09-27T18:36:11.324Z" }, + { url = "https://files.pythonhosted.org/packages/bc/36/23578f29e9e582a4d0278e009b38081dbe363c5e7165113fad546918a232/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:d2ee202e79d8ed691ceebae8e0486bd9a2cd4794cec4824e1c99b6f5009502f6", size = 20282, upload-time = "2025-09-27T18:36:12.573Z" }, + { url = "https://files.pythonhosted.org/packages/56/21/dca11354e756ebd03e036bd8ad58d6d7168c80ce1fe5e75218e4945cbab7/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:177b5253b2834fe3678cb4a5f0059808258584c559193998be2601324fdeafb1", size = 20745, upload-time = "2025-09-27T18:36:13.504Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2e/5898933336b61975ce9dc04decbc0a7f2fee78c30353c5efba7f2d6ff27a/markupsafe-3.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a", size = 12058, upload-time = "2025-09-27T18:36:19.444Z" }, + { url = "https://files.pythonhosted.org/packages/1d/09/adf2df3699d87d1d8184038df46a9c80d78c0148492323f4693df54e17bb/markupsafe-3.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50", size = 24287, upload-time = "2025-09-27T18:36:20.768Z" }, + { url = "https://files.pythonhosted.org/packages/30/ac/0273f6fcb5f42e314c6d8cd99effae6a5354604d461b8d392b5ec9530a54/markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf", size = 22940, upload-time = "2025-09-27T18:36:22.249Z" }, + { url = "https://files.pythonhosted.org/packages/19/ae/31c1be199ef767124c042c6c3e904da327a2f7f0cd63a0337e1eca2967a8/markupsafe-3.0.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f", size = 21887, upload-time = "2025-09-27T18:36:23.535Z" }, + { url = "https://files.pythonhosted.org/packages/b2/76/7edcab99d5349a4532a459e1fe64f0b0467a3365056ae550d3bcf3f79e1e/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a", size = 23692, upload-time = "2025-09-27T18:36:24.823Z" }, + { url = "https://files.pythonhosted.org/packages/a4/28/6e74cdd26d7514849143d69f0bf2399f929c37dc2b31e6829fd2045b2765/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115", size = 21471, upload-time = "2025-09-27T18:36:25.95Z" }, + { url = "https://files.pythonhosted.org/packages/62/7e/a145f36a5c2945673e590850a6f8014318d5577ed7e5920a4b3448e0865d/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a", size = 22923, upload-time = "2025-09-27T18:36:27.109Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "mlx" +version = "0.32.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mlx-metal", marker = "sys_platform != 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/6d/369081cd86b32237910837ec2e8ed97c9c468f5e10798b28ee273301cc4b/mlx-0.32.3-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:0512049c81fded99aa6af2d039ae1fccf113ff1a5f2df937a596578aabd46869", size = 655804, upload-time = "2026-09-29T00:14:41.597Z" }, + { url = "https://files.pythonhosted.org/packages/fb/91/d68f66bb29e8de96d255b784a105d8a2fc26f6e844b7920ff19ded517c4d/mlx-0.32.3-cp310-cp310-macosx_15_0_arm64.whl", hash = "sha256:b57996fd2b72207a796656bc39508f20089692815297120039284edb296b5fff", size = 655802, upload-time = "2026-09-29T00:14:43.321Z" }, + { url = "https://files.pythonhosted.org/packages/f7/72/8056c38934512ba9dcc63c0f1df5bdc67a27881e205fb51730fdc8832b17/mlx-0.32.3-cp310-cp310-macosx_26_0_arm64.whl", hash = "sha256:b14744a1b3a9dd48e06a62c448710a798ceda9451d7d37a4654a58b5adfa97f8", size = 655957, upload-time = "2026-09-29T00:14:45.001Z" }, + { url = "https://files.pythonhosted.org/packages/2f/0e/333d6d51386c940d5b364b815e6c1c24b26a565de8e04de38eb1433a16ec/mlx-0.32.3-cp310-cp310-manylinux_2_35_aarch64.whl", hash = "sha256:c0764ce1a7a7f19768d6b3473387dcac03cac8b0dc74b7e1c8d05b451ff4786c", size = 693228, upload-time = "2026-09-29T00:14:46.674Z" }, + { url = "https://files.pythonhosted.org/packages/1d/e6/12dd16dc0306980048395cefb13b6d5543f23db997730763f901fb632969/mlx-0.32.3-cp310-cp310-manylinux_2_35_x86_64.whl", hash = "sha256:d61942263d4603814622cd72bb5991b8a3cb59970bde8a7b40b897a7f088fb84", size = 760238, upload-time = "2026-09-29T00:14:48.049Z" }, + { url = "https://files.pythonhosted.org/packages/a6/ed/7ced6107f0eb4959659e2a3bed2e402b70eb117d3c8941d4c9b21e9c3ca8/mlx-0.32.3-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:8439baa207db769df1df82a722b0ee0c2362f75c23c0740b7472c3d279253bae", size = 655739, upload-time = "2026-09-29T00:14:51.18Z" }, + { url = "https://files.pythonhosted.org/packages/d8/98/44ed08071e8ab99800d01ba7d6683a1f565233bae000d68aa7fe3a9ff11f/mlx-0.32.3-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:d9ff05b868748b846a02f59913e9aa9d7b576827db92cc358ec1cf660f7f90a5", size = 655739, upload-time = "2026-09-29T00:14:52.703Z" }, + { url = "https://files.pythonhosted.org/packages/3c/26/893e65bed5616901b7e8554ddf278bfe81037e1efa67736a7ea0f6924891/mlx-0.32.3-cp311-cp311-macosx_26_0_arm64.whl", hash = "sha256:4e1a8fa78461c6cf1f9473e3bc8a3c4e362f1aaa613b435787f2844fad522bc7", size = 655786, upload-time = "2026-09-29T00:14:54.266Z" }, + { url = "https://files.pythonhosted.org/packages/1f/51/084a2aee6485e87fc148c2aa87265ea4668e7e174ce9de8ac00b13af9c5b/mlx-0.32.3-cp311-cp311-manylinux_2_35_aarch64.whl", hash = "sha256:46df0a5489edcc071cec3d9dd1f7894d492888c9d425effdc8831f4e7021108d", size = 693156, upload-time = "2026-09-29T00:14:56.066Z" }, + { url = "https://files.pythonhosted.org/packages/3d/7a/b277953da0570bd126d7efac397b79f76ad2ff2bad6f710e555791c7cb56/mlx-0.32.3-cp311-cp311-manylinux_2_35_x86_64.whl", hash = "sha256:a14387cd679018a32e802fdaa58b1e9845a278cacd309955eddca95757988ce0", size = 760589, upload-time = "2026-09-29T00:14:57.638Z" }, + { url = "https://files.pythonhosted.org/packages/26/8c/2f8b796885b4c02ab9644916bebf656cd1d6d366df94d174b046903ea250/mlx-0.32.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:7aecc3d4a904279492dae58e28378a3c4bec4533e7066e64ae5c78c7a5511f22", size = 651041, upload-time = "2026-09-29T00:22:36.481Z" }, + { url = "https://files.pythonhosted.org/packages/6e/06/aacd94119d4a999d43f07f24b02904d7f895ee074ed4678b28456c319ea3/mlx-0.32.3-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:a6241ea8f56a08a6c3a06bbaad7d7fc2fc1c05f23757c61fe0a8797325ac6d46", size = 651042, upload-time = "2026-09-29T00:22:38.071Z" }, + { url = "https://files.pythonhosted.org/packages/4a/c1/13ff85d72f7cf01239c1031688b11cf1c7a3f0c7380636728976e510bfbd/mlx-0.32.3-cp312-cp312-macosx_26_0_arm64.whl", hash = "sha256:72711cb23cc6dccc5a460f63f39c11931a00b2c6dcd6fba253921090072e21bc", size = 651118, upload-time = "2026-09-29T00:22:39.503Z" }, + { url = "https://files.pythonhosted.org/packages/0e/c7/68b17902b5c38667e1ea1303a0c430ba1375cc5eea5e40d6b3a48b4d9f6d/mlx-0.32.3-cp312-cp312-manylinux_2_35_aarch64.whl", hash = "sha256:e75fa55161d9477e694f852cad0d383c6b4a5c3b505a5d08339f0e2260c2556c", size = 680634, upload-time = "2026-09-29T00:22:41.394Z" }, + { url = "https://files.pythonhosted.org/packages/9a/15/b7d62225c126c73ef65cc3ea9478f4c06bb7a4b39cf9e12237b9ca88554a/mlx-0.32.3-cp312-cp312-manylinux_2_35_x86_64.whl", hash = "sha256:b91012fdd5ec5cbedcec2cae6adfaad9a8f547023e570fc5781cb7eb4bfe2567", size = 754620, upload-time = "2026-09-29T00:22:43.068Z" }, + { url = "https://files.pythonhosted.org/packages/8c/e1/0039c25cb2e40e605de6409e6ede52b01b7c79d6a0727e0c103ec60d88b2/mlx-0.32.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f5716b51fea3d06352d4d8c851fb6c5141b1a7de5a5164d44de5a6d368d84a7e", size = 651142, upload-time = "2026-09-29T00:22:48.119Z" }, + { url = "https://files.pythonhosted.org/packages/29/dc/01ac6c32fb405a542be80f021bc978125af51d2ee5ae1db624bfd954e756/mlx-0.32.3-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:656df5e36b9bfa862fd633d55dee4da072d292def9f27f22134c594e03ff28a7", size = 651143, upload-time = "2026-09-29T00:22:49.702Z" }, + { url = "https://files.pythonhosted.org/packages/53/a9/70bf746c2cd13a198a429e33244d2f072ecd0c89508299a636529d3317cb/mlx-0.32.3-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:75333afef55819afd2a31e87ffbed0bf59f379bca21771e6b13a95a675e4c85b", size = 651075, upload-time = "2026-09-29T00:22:51.636Z" }, + { url = "https://files.pythonhosted.org/packages/f7/2c/d9a53850cb64f1dec9852aea3951b5aa5f83e8bce6b27800dedfc4147022/mlx-0.32.3-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:097856504e6cf5944b9aaa7c2e1d5be0693428abd9175429cae3893e7ec2f681", size = 680699, upload-time = "2026-09-29T00:22:53.167Z" }, + { url = "https://files.pythonhosted.org/packages/e5/14/fe8cba4d2c35f746a2631bbeb2ebc344eccaa5599aee6f4f6ca4a7acb7a4/mlx-0.32.3-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:936cc82b0c050686c43230c85f3a0bd9f179f650377ce9590f7af3ba659e2fce", size = 754734, upload-time = "2026-09-29T00:22:54.991Z" }, + { url = "https://files.pythonhosted.org/packages/68/73/f1511b38c9b200fc84bff84470ff1e4303f2f7b47f614b14a2d3f4cd6226/mlx-0.32.3-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:65122a7eeb718df56641d7ac6f7fdb0ada573fd841d1fb0a61a6a443e33f11b5", size = 675637, upload-time = "2026-09-29T00:22:59.92Z" }, + { url = "https://files.pythonhosted.org/packages/3e/40/20733670ac7a5571e10d80fb510bc0fbed6d0917a8e34cdf6c3de52b5a8e/mlx-0.32.3-cp313-cp313t-macosx_15_0_arm64.whl", hash = "sha256:e13f5ce391293dab946edf18e7ae7bb55768beef1140deb509f67a2b14edcf7c", size = 675638, upload-time = "2026-09-29T00:23:01.529Z" }, + { url = "https://files.pythonhosted.org/packages/15/ac/65a1b9484f3eedb57dc73589af213d03821d13d99caa76430ad849be6114/mlx-0.32.3-cp313-cp313t-macosx_26_0_arm64.whl", hash = "sha256:15e8578c9d10b4fd4728243cd6c926bc46e1da267c6a849cdbceca21d51caa51", size = 675600, upload-time = "2026-09-29T00:23:03.249Z" }, + { url = "https://files.pythonhosted.org/packages/f7/53/91f4030bed4ce920ab781b4f3e8843d74bed5c171b425d3ecca8a5ae2830/mlx-0.32.3-cp313-cp313t-manylinux_2_35_aarch64.whl", hash = "sha256:942f0da5328986cdae29ec0eba57b2a08eb9dc6125cfe196853de7498f2c5359", size = 686475, upload-time = "2026-09-29T00:23:04.777Z" }, + { url = "https://files.pythonhosted.org/packages/c1/5c/52651813664c3869d9ea4aab93d027f30b7208554254106120256bd90c4f/mlx-0.32.3-cp313-cp313t-manylinux_2_35_x86_64.whl", hash = "sha256:19da79adad1aa9653c70819040ac58d65d7a9610a375a30d03c7cce708d4d2bb", size = 755078, upload-time = "2026-09-29T00:23:06.436Z" }, +] + +[package.optional-dependencies] +cpu = [ + { name = "mlx-cpu" }, +] + +[[package]] +name = "mlx-cpu" +version = "0.32.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/d5/6c7275a441674afc9d4de6a169cd7907c205e2ca4e47a34c826c589dd313/mlx_cpu-0.32.3-py3-none-manylinux_2_35_aarch64.whl", hash = "sha256:59223b7ce089f2bc2b63f520332f03d784bcdbe6301184a9e4c3afa0bf2eb9fc", size = 8837475, upload-time = "2026-09-29T00:14:23.445Z" }, + { url = "https://files.pythonhosted.org/packages/83/0a/8e9be0b5f1551e5c27280d73599cf15d90f48db8ef029dd21ad3de6f2ec4/mlx_cpu-0.32.3-py3-none-manylinux_2_35_x86_64.whl", hash = "sha256:b85e30e51b121387f735eb56a4b874d0a21565cc36fa2a88ec1f1000d1ae28e6", size = 10358218, upload-time = "2026-09-29T00:14:26.124Z" }, +] + +[[package]] +name = "mlx-lm" +version = "0.31.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jinja2" }, + { name = "mlx", marker = "sys_platform != 'linux'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "numpy", version = "2.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "protobuf" }, + { name = "pyyaml" }, + { name = "sentencepiece" }, + { name = "transformers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/84/94/9a38d6b0c6fcca995b9136c94eb7da1e9c5165652edf228b96b29960fa7a/mlx_lm-0.31.3.tar.gz", hash = "sha256:61eb0e3ba09444f77f874aff295401d7ccd20b39495cbbce0c782a15474ce733", size = 304318, upload-time = "2026-04-22T07:37:27.922Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/02/9a67b8e4f87e3e2e5cd7b1ad79304b93c09a0db6af34bee75e6551c06c60/mlx_lm-0.31.3-py3-none-any.whl", hash = "sha256:758cfddf1180053b7613db76fad3d246a331a2a905808e1164a275621fc983b8", size = 408890, upload-time = "2026-04-22T07:37:25.965Z" }, +] + +[[package]] +name = "mlx-metal" +version = "0.32.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/ed/8510985bcb24a6eae359a5441ce4cedcee13b84258a010cd0e515e9fb0bc/mlx_metal-0.32.3-py3-none-macosx_14_0_arm64.whl", hash = "sha256:fdba88ed9131448cf163be2c7ad6374746d8818205f28657aa91b6c69050f691", size = 42205858, upload-time = "2026-09-29T00:14:27.819Z" }, + { url = "https://files.pythonhosted.org/packages/3c/dc/4f852295fbe357bb9d3a2bfd55361a520b94ded873b72dc58b34a7f507cd/mlx_metal-0.32.3-py3-none-macosx_15_0_arm64.whl", hash = "sha256:feea99fe4d01f9e9f33ca73d869f988219a1725de1c18870baa7d5cd5aef6cb3", size = 42187333, upload-time = "2026-09-29T00:14:33.861Z" }, + { url = "https://files.pythonhosted.org/packages/ee/38/cb985ca86979ca6f954a2a45eb6e3f9b787c5d55227564a04f87953b3d56/mlx_metal-0.32.3-py3-none-macosx_26_0_arm64.whl", hash = "sha256:34ae9b83ad2f0ccdd3e5d48ec35176e7119f57069eef187122916dc941a4ae1f", size = 67661845, upload-time = "2026-09-29T00:14:41.248Z" }, +] + +[[package]] +name = "numpy" +version = "2.2.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11' and platform_machine == 'arm64' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform == 'linux'", +] +sdist = { url = "https://files.pythonhosted.org/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/c2/4b9221495b2a132cc9d2eb862e21d42a009f5a60e45fc44b00118c174bff/numpy-2.2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e41fd67c52b86603a91c1a505ebaef50b3314de0213461c7a6e99c9a3beff90", size = 14360048, upload-time = "2025-05-17T21:28:21.406Z" }, + { url = "https://files.pythonhosted.org/packages/fd/77/dc2fcfc66943c6410e2bf598062f5959372735ffda175b39906d54f02349/numpy-2.2.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:37e990a01ae6ec7fe7fa1c26c55ecb672dd98b19c3d0e1d1f326fa13cb38d163", size = 5340542, upload-time = "2025-05-17T21:28:30.931Z" }, + { url = "https://files.pythonhosted.org/packages/eb/17/96a3acd228cec142fcb8723bd3cc39c2a474f7dcf0a5d16731980bcafa95/numpy-2.2.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:efd28d4e9cd7d7a8d39074a4d44c63eda73401580c5c76acda2ce969e0a38e83", size = 14297320, upload-time = "2025-05-17T21:29:02.78Z" }, + { url = "https://files.pythonhosted.org/packages/b4/63/3de6a34ad7ad6646ac7d2f55ebc6ad439dbbf9c4370017c50cf403fb19b5/numpy-2.2.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc7b73d02efb0e18c000e9ad8b83480dfcd5dfd11065997ed4c6747470ae8915", size = 16801050, upload-time = "2025-05-17T21:29:27.675Z" }, + { url = "https://files.pythonhosted.org/packages/07/b6/89d837eddef52b3d0cec5c6ba0456c1bf1b9ef6a6672fc2b7873c3ec4e2e/numpy-2.2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:74d4531beb257d2c3f4b261bfb0fc09e0f9ebb8842d82a7b4209415896adc680", size = 15807034, upload-time = "2025-05-17T21:29:51.102Z" }, + { url = "https://files.pythonhosted.org/packages/01/c8/dc6ae86e3c61cfec1f178e5c9f7858584049b6093f843bca541f94120920/numpy-2.2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8fc377d995680230e83241d8a96def29f204b5782f371c532579b4f20607a289", size = 18614185, upload-time = "2025-05-17T21:30:18.703Z" }, + { url = "https://files.pythonhosted.org/packages/b3/2b/64e1affc7972decb74c9e29e5649fac940514910960ba25cd9af4488b66c/numpy-2.2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c820a93b0255bc360f53eca31a0e676fd1101f673dda8da93454a12e23fc5f7a", size = 14406743, upload-time = "2025-05-17T21:31:41.087Z" }, + { url = "https://files.pythonhosted.org/packages/4a/9f/0121e375000b5e50ffdd8b25bf78d8e1a5aa4cca3f185d41265198c7b834/numpy-2.2.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3d70692235e759f260c3d837193090014aebdf026dfd167834bcba43e30c2a42", size = 5352616, upload-time = "2025-05-17T21:31:50.072Z" }, + { url = "https://files.pythonhosted.org/packages/52/b8/7f0554d49b565d0171eab6e99001846882000883998e7b7d9f0d98b1f934/numpy-2.2.6-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b64d8d4d17135e00c8e346e0a738deb17e754230d7e0810ac5012750bbd85a5a", size = 14312005, upload-time = "2025-05-17T21:32:23.332Z" }, + { url = "https://files.pythonhosted.org/packages/b3/dd/2238b898e51bd6d389b7389ffb20d7f4c10066d80351187ec8e303a5a475/numpy-2.2.6-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba10f8411898fc418a521833e014a77d3ca01c15b0c6cdcce6a0d2897e6dbbdf", size = 16821570, upload-time = "2025-05-17T21:32:47.991Z" }, + { url = "https://files.pythonhosted.org/packages/83/6c/44d0325722cf644f191042bf47eedad61c1e6df2432ed65cbe28509d404e/numpy-2.2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bd48227a919f1bafbdda0583705e547892342c26fb127219d60a5c36882609d1", size = 15818548, upload-time = "2025-05-17T21:33:11.728Z" }, + { url = "https://files.pythonhosted.org/packages/ae/9d/81e8216030ce66be25279098789b665d49ff19eef08bfa8cb96d4957f422/numpy-2.2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9551a499bf125c1d4f9e250377c1ee2eddd02e01eac6644c080162c0c51778ab", size = 18620521, upload-time = "2025-05-17T21:33:39.139Z" }, + { url = "https://files.pythonhosted.org/packages/66/ee/560deadcdde6c2f90200450d5938f63a34b37e27ebff162810f716f6a230/numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c", size = 14119362, upload-time = "2025-05-17T21:35:01.241Z" }, + { url = "https://files.pythonhosted.org/packages/3c/65/4baa99f1c53b30adf0acd9a5519078871ddde8d2339dc5a7fde80d9d87da/numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3", size = 5084103, upload-time = "2025-05-17T21:35:10.622Z" }, + { url = "https://files.pythonhosted.org/packages/f8/35/8c80729f1ff76b3921d5c9487c7ac3de9b2a103b1cd05e905b3090513510/numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87", size = 14018462, upload-time = "2025-05-17T21:35:42.174Z" }, + { url = "https://files.pythonhosted.org/packages/8c/3d/1e1db36cfd41f895d266b103df00ca5b3cbe965184df824dec5c08c6b803/numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249", size = 16527618, upload-time = "2025-05-17T21:36:06.711Z" }, + { url = "https://files.pythonhosted.org/packages/61/c6/03ed30992602c85aa3cd95b9070a514f8b3c33e31124694438d88809ae36/numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49", size = 15505511, upload-time = "2025-05-17T21:36:29.965Z" }, + { url = "https://files.pythonhosted.org/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, + { url = "https://files.pythonhosted.org/packages/dc/9e/14520dc3dadf3c803473bd07e9b2bd1b69bc583cb2497b47000fed2fa92f/numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b", size = 14143006, upload-time = "2025-05-17T21:38:18.291Z" }, + { url = "https://files.pythonhosted.org/packages/4f/06/7e96c57d90bebdce9918412087fc22ca9851cceaf5567a45c1f404480e9e/numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d", size = 5076765, upload-time = "2025-05-17T21:38:27.319Z" }, + { url = "https://files.pythonhosted.org/packages/85/c5/e19c8f99d83fd377ec8c7e0cf627a8049746da54afc24ef0a0cb73d5dfb5/numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f", size = 14010719, upload-time = "2025-05-17T21:38:58.433Z" }, + { url = "https://files.pythonhosted.org/packages/19/49/4df9123aafa7b539317bf6d342cb6d227e49f7a35b99c287a6109b13dd93/numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f", size = 16526072, upload-time = "2025-05-17T21:39:22.638Z" }, + { url = "https://files.pythonhosted.org/packages/b2/6c/04b5f47f4f32f7c2b0e7260442a8cbcf8168b0e1a41ff1495da42f42a14f/numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868", size = 15503213, upload-time = "2025-05-17T21:39:45.865Z" }, + { url = "https://files.pythonhosted.org/packages/17/0a/5cd92e352c1307640d5b6fec1b2ffb06cd0dabe7d7b8227f97933d378422/numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d", size = 18316632, upload-time = "2025-05-17T21:40:13.331Z" }, + { url = "https://files.pythonhosted.org/packages/61/66/d2de6b291507517ff2e438e13ff7b1e2cdbdb7cb40b3ed475377aece69f9/numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda", size = 14225144, upload-time = "2025-05-17T21:41:05.695Z" }, + { url = "https://files.pythonhosted.org/packages/e4/25/480387655407ead912e28ba3a820bc69af9adf13bcbe40b299d454ec011f/numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40", size = 5200217, upload-time = "2025-05-17T21:41:15.903Z" }, + { url = "https://files.pythonhosted.org/packages/b7/30/172c2d5c4be71fdf476e9de553443cf8e25feddbe185e0bd88b096915bcc/numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f", size = 14077935, upload-time = "2025-05-17T21:41:49.738Z" }, + { url = "https://files.pythonhosted.org/packages/12/fb/9e743f8d4e4d3c710902cf87af3512082ae3d43b945d5d16563f26ec251d/numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa", size = 16600122, upload-time = "2025-05-17T21:42:14.046Z" }, + { url = "https://files.pythonhosted.org/packages/12/75/ee20da0e58d3a66f204f38916757e01e33a9737d0b22373b3eb5a27358f9/numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571", size = 15586143, upload-time = "2025-05-17T21:42:37.464Z" }, + { url = "https://files.pythonhosted.org/packages/76/95/bef5b37f29fc5e739947e9ce5179ad402875633308504a52d188302319c8/numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1", size = 18385260, upload-time = "2025-05-17T21:43:05.189Z" }, + { url = "https://files.pythonhosted.org/packages/af/30/feba75f143bdc868a1cc3f44ccfa6c4b9ec522b36458e738cd00f67b573f/numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543", size = 16643476, upload-time = "2025-05-17T21:45:11.871Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.11.*' and platform_machine == 'arm64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'linux'", +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/0d/f5957185c0ee2f3e12f78715aa9e3b353fd83633316c8532b38faa37e3f6/numpy-2.4.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:110f8b71aacb688ec69062bb7f6938a0f8acb01b7c1c4beb453c65b6d234584d", size = 14964111, upload-time = "2026-05-18T23:33:17.795Z" }, + { url = "https://files.pythonhosted.org/packages/ad/40/40a40ee0ddf7ceb782c49af278894b686e586d65d8c1889c8b5da01a3d7d/numpy-2.4.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:4cfe66903cc32a9921a6733d96b19bb6abf310397581bbad89c228f5abaf0ee8", size = 5469159, upload-time = "2026-05-18T23:33:20.654Z" }, + { url = "https://files.pythonhosted.org/packages/33/a8/6fa8c1a345a8c85dbb21932c447bee07c30a2c2a3f31e369c0a84b300147/numpy-2.4.6-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ab0a9c4ffb1a6d95ef519fe4247dba8eb6b18ad93999f76b7f657039acabd47", size = 15966692, upload-time = "2026-05-18T23:33:26.62Z" }, + { url = "https://files.pythonhosted.org/packages/02/03/74fe2a4cb3817d94d86402f2506554130a2f01414e299b5a843e5a8a957f/numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cd468399cfd2504718f0ba50e410dca55a170b61a02ad92bb18c8a65186e93", size = 16918164, upload-time = "2026-05-18T23:33:29.955Z" }, + { url = "https://files.pythonhosted.org/packages/c5/80/3615be3313f7e7696609bc194b9f0101da809df79e859bdb84e0cd043f46/numpy-2.4.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c2d37ab77531417474168eb79d6d80b14f821a966818505d03013d0833edb7a8", size = 17322877, upload-time = "2026-05-18T23:33:34.724Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ac/a691e0fe2675e370d0e08ff905adc49a1c8830e8cae03efe4477e92cd55d/numpy-2.4.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f407cb6b8e9d6d8c626bc73c945db1706035af8fd632295547bf1c9e46d092d6", size = 18651487, upload-time = "2026-05-18T23:33:38.217Z" }, + { url = "https://files.pythonhosted.org/packages/ea/12/92c4c131527599e8288d6918e888d88726f84d805d784b771f32408aeaef/numpy-2.4.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ebfb099f8dcf083deef3ac1ca4c1503f387cf76296fcb3816b66f5ecb5f54fdb", size = 14699246, upload-time = "2026-05-18T23:33:57.621Z" }, + { url = "https://files.pythonhosted.org/packages/ad/fe/c0a6b7b2ca128a8fb228575147073b660656734b8ebe4d76c8fd748dcc79/numpy-2.4.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3213d622a0283a39a93d188f3cf72b26862df52fbb4ca3697f51705016523d41", size = 5204410, upload-time = "2026-05-18T23:34:00.302Z" }, + { url = "https://files.pythonhosted.org/packages/c9/c6/50a46a6205feba2343f1d6d17438107c5dc491ed1c736e6ea68689fd906b/numpy-2.4.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f9fb9157b4ce2971008323afe46053787b526ef624fea915b261468a8421a0f", size = 15671012, upload-time = "2026-05-18T23:34:05.485Z" }, + { url = "https://files.pythonhosted.org/packages/99/60/14115e6364fa676c5397c2ad3004e527e9aa487abf5d0706ec81bbd08529/numpy-2.4.6-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f9849678c75fe7afa2d348ac842c168b0a4d3d61919687216dfc547976d853", size = 16645538, upload-time = "2026-05-18T23:34:09.265Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c5/693cbe59e57db94d2231fa519ca3978dc9e19da5a8f088588f5c6e947ff2/numpy-2.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c1a2af6c6ef86344a6b0db6b97834208bf598db514f2b155042439b62605601a", size = 17020706, upload-time = "2026-05-18T23:34:13.053Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fc/85b7c4eff9b4966ade25c2273cf7e7012e92366c032058653934b37de044/numpy-2.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e5805d5a22fd19c8ccff10a9561f9df94436b0545619ea579db2d3c35294bce2", size = 18368541, upload-time = "2026-05-18T23:34:17.024Z" }, + { url = "https://files.pythonhosted.org/packages/1b/30/a80189bcc7f5e4258b3fbc3968d909d1756f54d023299ecc39ad6fdb9ef8/numpy-2.4.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bf162abab1c1a736333192707cef898e735a5ca00f38f27eeedf44b39d9e85eb", size = 14693902, upload-time = "2026-05-18T23:34:33.013Z" }, + { url = "https://files.pythonhosted.org/packages/97/12/70b5d0d7c15e1ebb8a6a84a8caa1d19e181d84fb58bb6d70aca29099dec1/numpy-2.4.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:043191bfa8eab18c776647b62723ac9dddece59743b13f49b2016094129c2b3f", size = 5198992, upload-time = "2026-05-18T23:34:36.132Z" }, + { url = "https://files.pythonhosted.org/packages/bb/c5/7b863a97a91671a0338f4253bd3b5a3d3852f0692dae91711c9f4a10e787/numpy-2.4.6-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72fbe16c6fac95aedf5937fa873445cec2110be35d8a4e9433d7501fd98dae6b", size = 15669392, upload-time = "2026-05-18T23:34:41.257Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9d/3584b9984ca4c047aea75214ce1a4c4c73d849bd71b604264b7f5653f8a8/numpy-2.4.6-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a7830bab239b79cda9c08c2da014761cafb48da6150e1da17ac06283f43b6089", size = 16633220, upload-time = "2026-05-18T23:34:45.075Z" }, + { url = "https://files.pythonhosted.org/packages/05/ae/7c67fba23bd98caec7c99261f3a16072ade14813486b0282cb29846de832/numpy-2.4.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ef4aea96ce4d3b074422cb4f2f64e216bf9e213004bb58ecfdf50ea02ea8eb9a", size = 17020800, upload-time = "2026-05-18T23:34:49.065Z" }, + { url = "https://files.pythonhosted.org/packages/d9/5d/3b6725cb31d983c5e66916f5d36f6d7e5521129e4c4404d64f918292a5b6/numpy-2.4.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dfa20cc6ca228e6b155b11da03825975ce66aea520985dbbddf0f2a5a495c605", size = 18357600, upload-time = "2026-05-18T23:34:52.709Z" }, + { url = "https://files.pythonhosted.org/packages/63/6d/cc5619247c8f4204e507f5883528372e4ac4bb189e579fb859a12e480b1f/numpy-2.4.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:112b06a867b235ef466ed3508ddf0238050df9c727cafb5301ac385b899189a1", size = 14821197, upload-time = "2026-05-18T23:35:05.468Z" }, + { url = "https://files.pythonhosted.org/packages/00/58/f1c39161c87d9e9bed660f1ed4bafc0e403d5ec9650b6dd77aead07d489b/numpy-2.4.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:eaf7fa2de5c0be8ae6ff8e9bea2ccd725e980541244521d8d4b5f3354a27babe", size = 5326287, upload-time = "2026-05-18T23:35:08.693Z" }, + { url = "https://files.pythonhosted.org/packages/eb/0f/037e64c494b67581ae18193d770adef354c41f3f2c8ebf865602d949bf8f/numpy-2.4.6-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20", size = 15728070, upload-time = "2026-05-18T23:35:14.79Z" }, + { url = "https://files.pythonhosted.org/packages/21/a6/5d2bae9c9542eb4df16dc9c46dc79c186e9bad53805dfa5399a6023c6db0/numpy-2.4.6-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede83e07a75dd06bc501566c1eca2afc0d61677c1472ac9ad93fdee6e638a48d", size = 16681752, upload-time = "2026-05-18T23:35:18.836Z" }, + { url = "https://files.pythonhosted.org/packages/92/14/23d1dfb410ae362cd59ce53e936b1513d545eb40db3949ced632e19a459e/numpy-2.4.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:68bb27509ac1b9a3443094260f6326150663b06abe40b73a2f81160623da5b67", size = 17086024, upload-time = "2026-05-18T23:35:22.52Z" }, + { url = "https://files.pythonhosted.org/packages/4b/6e/23595a2c642cdf3bc567877064bdd7f91c8b0038a4453cf2daf7248eafe9/numpy-2.4.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a0df0043bdb289bde1f62da130d20df23d58b45429f752bc7a8fc5325a225ecd", size = 18403398, upload-time = "2026-05-18T23:35:26.398Z" }, + { url = "https://files.pythonhosted.org/packages/44/53/f481bef68011740f8849418d82db07230e825013f31f4eef5ba5b805316a/numpy-2.4.6-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d6da64deb6b8ed903e7560180a92f2d804ee1ba5eeb849ac2748b8c1aba1f6d7", size = 14889064, upload-time = "2026-05-18T23:36:53.879Z" }, + { url = "https://files.pythonhosted.org/packages/7f/57/42ed575c10ced8af951d426bc4e1f8aff16fd851db33f067036215a7f860/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:68a5124b13fa6cc2086764a20005d30bc0548146f7f5322f02fce212ca14317f", size = 5394157, upload-time = "2026-05-18T23:36:57.194Z" }, + { url = "https://files.pythonhosted.org/packages/1a/9c/c531f2293b91265d8b48e9b329f54fdd7ffae73cb4134ea10cca4237e9cc/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5dbbdb29840ca3d91ee0fece42fc29278886d908280bfec0a5846c6f901a3eb0", size = 15798374, upload-time = "2026-05-18T23:37:02.674Z" }, + { url = "https://files.pythonhosted.org/packages/1a/b0/413077f6b1153ed3cba361401c6783bbad6114804a000cc22eb71c13e190/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ad03c0965fb3c692200e74d458ca28c1dbb4ce96f9a479a8aa041ad5fabca02", size = 16747286, upload-time = "2026-05-18T23:37:06.327Z" }, +] + +[[package]] +name = "numpy" +version = "2.5.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.12' and platform_machine == 'arm64' and sys_platform == 'darwin'", + "python_full_version >= '3.12' and sys_platform == 'linux'", +] +sdist = { url = "https://files.pythonhosted.org/packages/13/01/11703282db468b85f6f7b8c7f22d058de5970d5c7e60a3a8aaa313c3de36/numpy-2.5.3.tar.gz", hash = "sha256:df2d5874ff183595a4ba404edd04f6bd9b5505c1d7708573f6a6c17489a67563", size = 20791231, upload-time = "2026-09-06T16:27:47.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/60/39/789131c1188c078dcb3a1692e72e1e050c68b88ffe72c9ccaac9bcd7a9cd/numpy-2.5.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f59a878c33d6b88122d80d239bb3b845d58708750b0cb06a09aebb9b18ec696c", size = 12009327, upload-time = "2026-09-06T16:24:32.491Z" }, + { url = "https://files.pythonhosted.org/packages/9c/59/a312e95696e5f601914dd8b6dd844692ba61670807417e24b68e337b5c70/numpy-2.5.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:a72f874bc9e10e4b8f80426fb49716d5141f64442a0c8418065093ec8017fbb0", size = 5445405, upload-time = "2026-09-06T16:24:35.071Z" }, + { url = "https://files.pythonhosted.org/packages/f1/32/84146fc020ad3c25f805f70ab60da46fe3c540a21369754a7e4369754b6f/numpy-2.5.3-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:76c2c1e6bfa5c84adc6434dfbf013aa92096a7985221762c8f11fedfd20fff58", size = 15687872, upload-time = "2026-09-06T16:24:39.751Z" }, + { url = "https://files.pythonhosted.org/packages/65/af/aa78d1a88805456e212b65461354cd943197fb9acecc4c90fd12295123a3/numpy-2.5.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b7e18c623bb5c95acb3b3328861272816ba199fb531921c5d6d0b675f1fde9e3", size = 16717410, upload-time = "2026-09-06T16:24:42.745Z" }, + { url = "https://files.pythonhosted.org/packages/3b/24/faa79d865e69a97ba17473b23a1b74094b2259c03e820c70297293b9ea49/numpy-2.5.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4f8929ee6c96bfbd7b4ed2032e0c03af86fe1826740ab61ddabf9072d06e57ff", size = 17040975, upload-time = "2026-09-06T16:24:45.961Z" }, + { url = "https://files.pythonhosted.org/packages/62/4a/8877e629445a7176297dffcaf9c485faa96a95d81728a62521ad55bd4c0f/numpy-2.5.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b5d93cf48f687479941d12b69c873ad2cc76bbd487f0091c2200636497f34034", size = 18476479, upload-time = "2026-09-06T16:24:49.35Z" }, + { url = "https://files.pythonhosted.org/packages/2f/06/9dc9e48b5e5e941c8b10350c5ff2d721da42a20517d911d15544246775ff/numpy-2.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:92f30e89b8ee0ecf363033576c422b2f58fed6a80bed0aa48dff6d14c654663e", size = 12003676, upload-time = "2026-09-06T16:25:03.475Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2a/98282aa5b8f58b1157d440bb6282eed47e3632a5de53a714fbab17e659fe/numpy-2.5.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f9a2353b37a1a9e78fd82b27ad7e2a32a2d036604d18f02b05e3136c62ca3b09", size = 5439695, upload-time = "2026-09-06T16:25:05.978Z" }, + { url = "https://files.pythonhosted.org/packages/73/85/735720d04ec197c5dcfacdfc9922667c7f1f5f496a279b7ba4d7c74c4cc7/numpy-2.5.3-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c76d5dde9f445058f83d0c02af00557a4db91de9a9a57c0df87d1535001d654b", size = 15681750, upload-time = "2026-09-06T16:25:11.173Z" }, + { url = "https://files.pythonhosted.org/packages/3a/1b/3b16a9bc514a440a7a0883684111dcb1ef1aee960af2ca95da8fc775f124/numpy-2.5.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a5fa86b80fd24bcd1aff83ad23be44ea323de3f787be8f8b15d4a65621e25321", size = 16708577, upload-time = "2026-09-06T16:25:14.171Z" }, + { url = "https://files.pythonhosted.org/packages/69/c4/386f397831b07328b639c96c5b62719346cf4baf07c68d927239752b1534/numpy-2.5.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bd4cb9ad3c7889b9b3fe0a9a9fb5d2ed26f9879bff2608d9f01aed147a20d231", size = 17042047, upload-time = "2026-09-06T16:25:17.582Z" }, + { url = "https://files.pythonhosted.org/packages/5f/3e/a700ecbf36e85ae8328fd3b0e12eeddc22ed6358a64cb2bd913e0d195d65/numpy-2.5.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1302b90c0e52281681b2975adfe8a860cb7b12216a27b4b0b4207c44bf7bccf0", size = 18465724, upload-time = "2026-09-06T16:25:20.949Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "protobuf" +version = "7.36.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/89/5b8517baa72f84a67b8a307ba953c91057af618bf40bf676f3c03551f8f0/protobuf-7.36.2.tar.gz", hash = "sha256:497d0463ff3316681da6c0b9e8d06cb465d61abce00b613ab42226175644d1bb", size = 512737, upload-time = "2026-09-17T20:07:59.326Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/72/98342feb672507c8f3a69e34b4fa8961f608edba5c1a48a6f47156d92cb5/protobuf-7.36.2-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:cbc70b17ee27e28894c7fee8bb04be1abead49e936bc70eb60052531eee2079e", size = 456039, upload-time = "2026-09-17T20:07:51.542Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ea/91fdf7c2b8bbd49cde056f00a9df6773532987e1c00fe2830b895af95c7e/protobuf-7.36.2-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:e11e1f0180583a2af89db6a2ecd9e8dc40aa6d2988ca175bfd0e6d12ea72d74e", size = 344219, upload-time = "2026-09-17T20:07:52.914Z" }, + { url = "https://files.pythonhosted.org/packages/17/ab/5fd5f8ece73fad885c5a09aa849b32d70472f954ba3a92d3bb5974ea953b/protobuf-7.36.2-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:f4fee11ec330d238b34a05c9b675f693c20415d1c5bd7d5320cc2f8a798eb9cf", size = 357223, upload-time = "2026-09-17T20:07:53.985Z" }, + { url = "https://files.pythonhosted.org/packages/db/f3/3996583dd2906297a637af12114deddf7658af6e683fedb83be061983fb5/protobuf-7.36.2-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:89f23aa53c24553a2416fd4fd1ec06f74fa42b14b546d8883128813f775bbfd2", size = 343223, upload-time = "2026-09-17T20:07:54.931Z" }, + { url = "https://files.pythonhosted.org/packages/e4/04/d52c7016b04b6c5108f26691f9d33ec82a9b65d041f1a9c771137693d618/protobuf-7.36.2-py3-none-any.whl", hash = "sha256:bdb3a345d48db958e6ce1f18e508beb0cc981d64f24088427549c866cd039f1e", size = 179806, upload-time = "2026-09-17T20:07:58.211Z" }, +] + +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/14/52d505b5c59ce73244f59c7a50ecf47093ce4765f116cdb98286a71eeca2/pyyaml-6.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:02ea2dfa234451bbb8772601d7b8e426c2bfa197136796224e50e35a78777956", size = 174019, upload-time = "2025-09-25T21:31:47.706Z" }, + { url = "https://files.pythonhosted.org/packages/43/f7/0e6a5ae5599c838c696adb4e6330a59f463265bfa1e116cfd1fbb0abaaae/pyyaml-6.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b30236e45cf30d2b8e7b3e85881719e98507abed1011bf463a8fa23e9c3e98a8", size = 740646, upload-time = "2025-09-25T21:31:49.21Z" }, + { url = "https://files.pythonhosted.org/packages/2f/3a/61b9db1d28f00f8fd0ae760459a5c4bf1b941baf714e207b6eb0657d2578/pyyaml-6.0.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:66291b10affd76d76f54fad28e22e51719ef9ba22b29e1d7d03d6777a9174198", size = 840793, upload-time = "2025-09-25T21:31:50.735Z" }, + { url = "https://files.pythonhosted.org/packages/7a/1e/7acc4f0e74c4b3d9531e24739e0ab832a5edf40e64fbae1a9c01941cabd7/pyyaml-6.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9c7708761fccb9397fe64bbc0395abcae8c4bf7b0eac081e12b809bf47700d0b", size = 770293, upload-time = "2025-09-25T21:31:51.828Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ef/abd085f06853af0cd59fa5f913d61a8eab65d7639ff2a658d18a25d6a89d/pyyaml-6.0.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:418cf3f2111bc80e0933b2cd8cd04f286338bb88bdc7bc8e6dd775ebde60b5e0", size = 732872, upload-time = "2025-09-25T21:31:53.282Z" }, + { url = "https://files.pythonhosted.org/packages/1f/15/2bc9c8faf6450a8b3c9fc5448ed869c599c0a74ba2669772b1f3a0040180/pyyaml-6.0.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5e0b74767e5f8c593e8c9b5912019159ed0533c70051e9cce3e8b6aa699fcd69", size = 758828, upload-time = "2025-09-25T21:31:54.807Z" }, + { url = "https://files.pythonhosted.org/packages/16/19/13de8e4377ed53079ee996e1ab0a9c33ec2faf808a4647b7b4c0d46dd239/pyyaml-6.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:652cb6edd41e718550aad172851962662ff2681490a8a711af6a4d288dd96824", size = 175577, upload-time = "2025-09-25T21:32:00.088Z" }, + { url = "https://files.pythonhosted.org/packages/0c/62/d2eb46264d4b157dae1275b573017abec435397aa59cbcdab6fc978a8af4/pyyaml-6.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:10892704fc220243f5305762e276552a0395f7beb4dbf9b14ec8fd43b57f126c", size = 775556, upload-time = "2025-09-25T21:32:01.31Z" }, + { url = "https://files.pythonhosted.org/packages/10/cb/16c3f2cf3266edd25aaa00d6c4350381c8b012ed6f5276675b9eba8d9ff4/pyyaml-6.0.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:850774a7879607d3a6f50d36d04f00ee69e7fc816450e5f7e58d7f17f1ae5c00", size = 882114, upload-time = "2025-09-25T21:32:03.376Z" }, + { url = "https://files.pythonhosted.org/packages/71/60/917329f640924b18ff085ab889a11c763e0b573da888e8404ff486657602/pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b8bb0864c5a28024fac8a632c443c87c5aa6f215c0b126c449ae1a150412f31d", size = 806638, upload-time = "2025-09-25T21:32:04.553Z" }, + { url = "https://files.pythonhosted.org/packages/dd/6f/529b0f316a9fd167281a6c3826b5583e6192dba792dd55e3203d3f8e655a/pyyaml-6.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1d37d57ad971609cf3c53ba6a7e365e40660e3be0e5175fa9f2365a379d6095a", size = 767463, upload-time = "2025-09-25T21:32:06.152Z" }, + { url = "https://files.pythonhosted.org/packages/f2/6a/b627b4e0c1dd03718543519ffb2f1deea4a1e6d42fbab8021936a4d22589/pyyaml-6.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:37503bfbfc9d2c40b344d06b2199cf0e96e97957ab1c1b546fd4f87e53e5d3e4", size = 794986, upload-time = "2025-09-25T21:32:07.367Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, +] + +[[package]] +name = "regex" +version = "2026.9.29" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fc/f2/af1da9d3ceed77bfcdce40427d49ba0be94e4fe84245e3bfef68c10e75b6/regex-2026.9.29.tar.gz", hash = "sha256:8b5fcc4771732191b2b7d1dd68d8f0353f47f8d90b6150f6dce58bf1112442cb", size = 419199, upload-time = "2026-09-29T00:49:58.298Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/37/404442e7296554792e2de4c0c8f8b95ee491a3c227d6ff90d4cebf2dbde3/regex-2026.9.29-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:9916fda742cd4eede63b286f58c06718324265d727ce0856eb1aac86d0d150d6", size = 495402, upload-time = "2026-09-29T00:45:50.322Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1d/aefd12fb5cd62b748aa6ed28725b375c71102853fcbf15fb6481f806a2ba/regex-2026.9.29-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1d9fe8091b2e89d470df68a9331111ed008ae8aae6bf1e8e1fba4086a495c84e", size = 292130, upload-time = "2026-09-29T00:45:53.017Z" }, + { url = "https://files.pythonhosted.org/packages/7c/98/41fe60ae6ccf3e166fe7bc73495aa15e9a1a7fe3a1f9f0f95a3d64a2df0f/regex-2026.9.29-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fb00027a09a8f9f08028b40dce4c933cf73e4833240ed356583fdc9cfa721566", size = 793080, upload-time = "2026-09-29T00:45:54.331Z" }, + { url = "https://files.pythonhosted.org/packages/45/0d/944b13e8286ecba61a3baeca28b375550821a279800bdb18fac41ade6637/regex-2026.9.29-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:14e953ff3607c92d7675bf79c4d4509ef6782aa8c08509f179f9b3d6d0679e86", size = 865241, upload-time = "2026-09-29T00:45:56.065Z" }, + { url = "https://files.pythonhosted.org/packages/2a/b4/e1c32fdee0462f1373c52a73325cafbde653fac0b0cc1e9badac74154f4a/regex-2026.9.29-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0476e5bcbe6e1ba3d1c4cc7bbb1c3ba78e3b979b5c8a88d0a6a8cdd4992b8c84", size = 905496, upload-time = "2026-09-29T00:45:57.491Z" }, + { url = "https://files.pythonhosted.org/packages/ec/99/eae371ca63f7ea1f1eac025021527adb57ce3a4b67ff608b896fd31ea447/regex-2026.9.29-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4fb41211d2333eb930a51e0546a65999761cf1f572a4da56ef9b8a62966c06f2", size = 797643, upload-time = "2026-09-29T00:45:58.746Z" }, + { url = "https://files.pythonhosted.org/packages/dc/26/9e4f158459e37e41633c96f49766a465f801c5daba9507260609ed76d783/regex-2026.9.29-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:edf06545875f3efa31560d94121e95c7fd70d98b1dfedc0157097d79b13b52ea", size = 797591, upload-time = "2026-09-29T00:46:00.2Z" }, + { url = "https://files.pythonhosted.org/packages/86/60/5e1c5c6e85132e757ae77bd7693f1809046edb2c8853658c9200f4c8593c/regex-2026.9.29-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:6398d5145689503412cc1748895242598d8846b8967b851133b20dc2ed1e21e8", size = 774874, upload-time = "2026-09-29T00:46:01.491Z" }, + { url = "https://files.pythonhosted.org/packages/e5/13/301b394f32a8db86ae3280533c9186229595fed766376f18e87959f92f96/regex-2026.9.29-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:45010bcfe66df41522d56c9b6114e87ecc597a08970ff6a2ced24415c141ae5f", size = 780168, upload-time = "2026-09-29T00:46:02.817Z" }, + { url = "https://files.pythonhosted.org/packages/f6/27/0c647127db32760b2dfc4c38da9ca509a3c4a8d3d93a586121a10fef2701/regex-2026.9.29-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:a5758353650079898dc1b2b0e95aa51fa23a30d020e06f62c430dd08ee56cdd8", size = 856604, upload-time = "2026-09-29T00:46:04.152Z" }, + { url = "https://files.pythonhosted.org/packages/4f/18/cac15fb2829c33218ececf76333b66869a96e4f973480ec32494921750f0/regex-2026.9.29-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:6f7121a8914ed13fcfe2099f895341bfb789f004d4c5a0bdece8fa667da10849", size = 760602, upload-time = "2026-09-29T00:46:05.794Z" }, + { url = "https://files.pythonhosted.org/packages/b8/55/967676e4089a25d5f54b9d62522f9fe09b87677d1db58f207888d4ccfbe3/regex-2026.9.29-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:b9d74e4eee9ddb64c2e92d5d61472c59c21684c059eb7b68767be9628e977859", size = 846830, upload-time = "2026-09-29T00:46:07.258Z" }, + { url = "https://files.pythonhosted.org/packages/f1/8d/bb4165d61860ed443da5fb0fafe629c1ee9202526394eead7925bdee2be4/regex-2026.9.29-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:143533cc4b6fbc5b95aca0a5b8d541088d374831593def000ec89322c220221d", size = 785445, upload-time = "2026-09-29T00:46:08.89Z" }, + { url = "https://files.pythonhosted.org/packages/e8/6b/6dea87689c3a06a6e79d254bf824e6f3e3d724b5ba027c6112559aa6cd2c/regex-2026.9.29-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6abb75ab16bc3281714a5b99548a2225db70dba1f995f6d7f7419b76eb5a8fbe", size = 495413, upload-time = "2026-09-29T00:46:14.51Z" }, + { url = "https://files.pythonhosted.org/packages/b1/07/9bf3607d8d13a12e436ab9d63f9791e10706827d535695b23964ad79fd79/regex-2026.9.29-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:066d0e3dbfdd739bce2bf8c2a41dd16f73e3d8adc2eb06dd803a36a307f56075", size = 292123, upload-time = "2026-09-29T00:46:17.646Z" }, + { url = "https://files.pythonhosted.org/packages/64/6b/32c2e6fc617e1d3f247e250fea31a9a35b1265bd32f585968aa13b9999b9/regex-2026.9.29-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7020ed44df30b3aa492c00ee3b52d0548c1f30c2c6c5bb13ae897680900d3413", size = 799537, upload-time = "2026-09-29T00:46:18.976Z" }, + { url = "https://files.pythonhosted.org/packages/bf/72/f041177f3c7a4606f7c81a95fe7eea03e2a0c4e8bff9e439a01432cbc9f2/regex-2026.9.29-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ae4613d7d9dda60fcba95f846cc6f808017f1843f392cf9daad14a6534493d71", size = 873255, upload-time = "2026-09-29T00:46:20.684Z" }, + { url = "https://files.pythonhosted.org/packages/d0/4e/a78948e11dd715e0e46716c2e0f3404b3fe6a44e2a2e9abdc7d965cab2b3/regex-2026.9.29-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bec37990e3d6121f29ecfb594bd8f1bf009e9f7926daba2e50e3b27d3892a783", size = 914570, upload-time = "2026-09-29T00:46:22.599Z" }, + { url = "https://files.pythonhosted.org/packages/8a/70/aa08d1d2b294894b365e5f8ba5380fe3f8546acdb81f10639dfd74209c37/regex-2026.9.29-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:612b709381c0355b70d89cdb51b7f670591ed5cbbc0e3b5337488019dc667b65", size = 804978, upload-time = "2026-09-29T00:46:23.981Z" }, + { url = "https://files.pythonhosted.org/packages/21/32/1b03534c4715aca3b564416d28d518083ed4dab3bc913267600d2256140d/regex-2026.9.29-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a760da040b47767b4b873adfb7c3b691e9ba2fc60f113f9d0b88f1a62f323e85", size = 779297, upload-time = "2026-09-29T00:46:25.318Z" }, + { url = "https://files.pythonhosted.org/packages/76/a7/378f6f558d9e4444af315a307c5953565a511d1e3666f1bb7bdc82012b6b/regex-2026.9.29-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:49ee178ca31c94621294bf9b8b676a92a2e6bba8af0529591753719e57edb621", size = 786961, upload-time = "2026-09-29T00:46:26.963Z" }, + { url = "https://files.pythonhosted.org/packages/59/13/79f0b1846f5f342f92ddbd4b27b18bcb86da96d902c1a0be26520bde98d7/regex-2026.9.29-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:5eeb8edc6110d9194a4d0d54610f64c37a31c605b5dbb7e407fc6ec7fa34a4a1", size = 863234, upload-time = "2026-09-29T00:46:28.58Z" }, + { url = "https://files.pythonhosted.org/packages/97/19/05af70dec9f2eed6ba34e08d2dcc6a48e7ae5e307659d5fe4201a5d7bbee/regex-2026.9.29-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:ccb64d887a9db1cd76dbc0f92051a1a478a2a67e7f56c62d915cb881d7734704", size = 766487, upload-time = "2026-09-29T00:46:29.941Z" }, + { url = "https://files.pythonhosted.org/packages/01/e1/9c7486d4afe8fdd1fe0ad60139f8aa91427381f409af6a29b609d8fdcb3a/regex-2026.9.29-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:9e4482589065c8ecd761cff522dcd85f2d39e62f551e37e025d1c7d54772def3", size = 854792, upload-time = "2026-09-29T00:46:31.358Z" }, + { url = "https://files.pythonhosted.org/packages/26/c7/49d008ff5f741d9a9799d7315556f3a12b983ff0fcd2cdfb62904bedafbf/regex-2026.9.29-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d60030baaa7bfbb02d650c126cdcddcb6e33dbff14d819434c8fa2fdcaeeeba5", size = 793050, upload-time = "2026-09-29T00:46:32.775Z" }, + { url = "https://files.pythonhosted.org/packages/84/48/3fdcde9a0baa84d7d25571223265d6e434e114763b438601d54a8028bf3e/regex-2026.9.29-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:dc79d36d0618752265f0d575915bdc5c5130ecb9c9f6b3bcefeae32e4bdfafcf", size = 497903, upload-time = "2026-09-29T00:46:38.938Z" }, + { url = "https://files.pythonhosted.org/packages/37/14/f3f0ba083d2094392d5eabf56db5ea6ba469fd6e927afd187042054ea68a/regex-2026.9.29-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f57dc6b8fef170f105d2cf5cdce254f47b137d7755086cf7050f47e16582abba", size = 293633, upload-time = "2026-09-29T00:46:41.959Z" }, + { url = "https://files.pythonhosted.org/packages/c9/72/67e7a8ce17f1aea49df215564048efb49cc8c2b31a0e0fc30f36838f8516/regex-2026.9.29-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f93bc1c3486ef3747e07c9d7c1d0a147b8fbaab975f80e348aed6f71309dfaca", size = 805885, upload-time = "2026-09-29T00:46:43.373Z" }, + { url = "https://files.pythonhosted.org/packages/f6/78/25436bcfd4d2260b4b4090094d55d7ab53ec8a1ab4865a0b8bcb33c7d5c0/regex-2026.9.29-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9e1d3a4cb7993b708f0ada8d0c84590efd853f169e7147d2202c9da503180242", size = 878344, upload-time = "2026-09-29T00:46:45.328Z" }, + { url = "https://files.pythonhosted.org/packages/97/e6/a09ec3a23ae41d6179880e67f0aace9284b2d95f2d7b326eff203f8eec5e/regex-2026.9.29-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:dabee8f4935e731fb46b2a3091bdda0d3d94b3bbfb907d2b4f12eefce4009619", size = 919181, upload-time = "2026-09-29T00:46:47.041Z" }, + { url = "https://files.pythonhosted.org/packages/26/83/d2fbd2e4e3afb1167daa825187d196f313cbaa1a4768f311fb041bb0e3d2/regex-2026.9.29-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:39ab5894d971f9ac68baa6eca5c50387db579cfcacf36ae8df3feceb1815e6d0", size = 807783, upload-time = "2026-09-29T00:46:48.894Z" }, + { url = "https://files.pythonhosted.org/packages/46/0b/eb429a7016610d44fc89a597163f8c9127505f0d7dc724dc9effbb6a3ac0/regex-2026.9.29-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c1a9a6651197fbed6f0212591418b9def774fc3f8324f78d1bf0e6a63e5f8aa1", size = 783465, upload-time = "2026-09-29T00:46:50.64Z" }, + { url = "https://files.pythonhosted.org/packages/1b/07/58a3c0153c7476898430f6a7cf3d9062a1d17fbea4f43399ecaf411c7b4c/regex-2026.9.29-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:87fb80cbe3557e27e7b28b995c2b2eedf689b8886f941ab93e0e288f0976518a", size = 793519, upload-time = "2026-09-29T00:46:52.396Z" }, + { url = "https://files.pythonhosted.org/packages/2a/e8/161b94d39164520e21a7befe0245569bf7fda4c7cf1fc4e2df2b5def49da/regex-2026.9.29-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3c5c2ef13797466aa64170cbb66ad98a32351dd4127694cea7199f80f213750d", size = 869293, upload-time = "2026-09-29T00:46:54.128Z" }, + { url = "https://files.pythonhosted.org/packages/8f/07/3b02ed829aa2decdc1955d222bd1e2f99d1c8bb4873bbb9a66b2f0a36bff/regex-2026.9.29-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:59b49507f47479e299a9e1bc41b5cb83a7afda0540625f1dbae886615978acbf", size = 770239, upload-time = "2026-09-29T00:46:56.106Z" }, + { url = "https://files.pythonhosted.org/packages/42/5b/ba61f6fe062eb8562e742367d177bb75370434138ef6c9d2a27114f8d613/regex-2026.9.29-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:0dd8af32e9f7b56b7f95cc1fd79b23054c3bdc172392ae560acc24d57b7ffe71", size = 861973, upload-time = "2026-09-29T00:46:57.665Z" }, + { url = "https://files.pythonhosted.org/packages/cc/27/767259b20e8a842948990f5e99138d6c077248fd42f8b5468b1d9ca4b814/regex-2026.9.29-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:db5e82ba15c142425b8406690032df89e39cca4a2e8afbbb9a3d84edc2373ac3", size = 796470, upload-time = "2026-09-29T00:46:59.236Z" }, + { url = "https://files.pythonhosted.org/packages/79/d5/6080f7d1a6e7e36aa720f806ac93c035ba39c209ae6cc510e8ef4c0279c6/regex-2026.9.29-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f1a0d5117230dd46b399a30a38afa44f79c99f3168988fdc4f425c3f928b39df", size = 497622, upload-time = "2026-09-29T00:47:08.251Z" }, + { url = "https://files.pythonhosted.org/packages/11/9e/aa0f4cde3bc4688c1d58b0cd8415edd708339bc0bc401a195b0b1e8c8f0c/regex-2026.9.29-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c90fcf7804ea0a54b896ce0f2b9565350220b8d4890fd0db461a476a4c687963", size = 293472, upload-time = "2026-09-29T00:47:11.723Z" }, + { url = "https://files.pythonhosted.org/packages/90/d4/e835c487850ed922a8d6074f953b888c8ea99775c76b9ed5f8a4d72eab92/regex-2026.9.29-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e11edba5bc344a32b029a7af9d4b3173982dd79eeafa0b9dbd787364414b0509", size = 805966, upload-time = "2026-09-29T00:47:13.235Z" }, + { url = "https://files.pythonhosted.org/packages/2c/57/ba8809847fbae8d2cbc71367c6ded510a7ec88bf52493c65efc1acf4effb/regex-2026.9.29-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bb90e7177944b6684738c1fc36aabd2dd00d1de3be7dbe09f91e196f1bc0dc81", size = 878346, upload-time = "2026-09-29T00:47:14.877Z" }, + { url = "https://files.pythonhosted.org/packages/1a/52/e3da19fc3cc15ef67ab67e121e87887c3bccfdb683a7a9ec557c460ca5b7/regex-2026.9.29-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d06fcdecc10fc7954d7c8f27a03c96055fe525274dc84a7b0dbdc3d6b9e03dab", size = 919250, upload-time = "2026-09-29T00:47:16.622Z" }, + { url = "https://files.pythonhosted.org/packages/9a/8e/c1ed81f55f992f6aa0b699a592a50c1ce9e6d44ff1aee2c14c0537dcef9c/regex-2026.9.29-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d49c18f1ea294cf4adde2e5ac256e98c82ea9d708462ce4bf799dffa7cfe8a2c", size = 807902, upload-time = "2026-09-29T00:47:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/ad/bc/5a6886eb470e41040e21e05b75024a18b6ebfe7ea400b72094a60f949101/regex-2026.9.29-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3e778bfccd63075167709136afbc251c1f683758d5bf49c803c60ac3f894ce6b", size = 783514, upload-time = "2026-09-29T00:47:19.921Z" }, + { url = "https://files.pythonhosted.org/packages/cb/52/6d951d453b023c6edb880f1ba474291b53b8ce1cc438b96a9db6d791d991/regex-2026.9.29-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:686ac5350fceae63830bb98805fcb8039325bf4c06d9f6f048ff65229d5bffa5", size = 793466, upload-time = "2026-09-29T00:47:21.552Z" }, + { url = "https://files.pythonhosted.org/packages/99/b9/d5a41adc08360f5eee0dc4846c578f002366947211fc8af5a69a64ee7b9f/regex-2026.9.29-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:26ec4ccce55aa533fbd603d08911b01101a8fcfec987845ac3ae2c7087b2bde3", size = 869464, upload-time = "2026-09-29T00:47:23.276Z" }, + { url = "https://files.pythonhosted.org/packages/4b/32/d76c9d91f5d798e2e9e67f6f85ec4ae35445ac425f7454797311cecb80ca/regex-2026.9.29-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:a655d34b2a6943af32401f3d94f72e9d731f6ad16285815550bf2b4ee69d420a", size = 770278, upload-time = "2026-09-29T00:47:25.193Z" }, + { url = "https://files.pythonhosted.org/packages/24/00/aeebdb540c620a0f7317f6d6fad80a47729ecf0599a24b5c34ec155351f5/regex-2026.9.29-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:0c992c19cd45058a4b92f68f139c93db168b48fb1f322c9a7cd620806afb6b51", size = 861949, upload-time = "2026-09-29T00:47:27.005Z" }, + { url = "https://files.pythonhosted.org/packages/12/62/d0314bcedfd3586197e4596931fa220260eb2385bf53184e5b9ae67db24b/regex-2026.9.29-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ebb8912f565b8cdbbf27debfe00df04202c20e2f651b9e32767930c5eace3621", size = 796602, upload-time = "2026-09-29T00:47:29.233Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, +] + +[[package]] +name = "sentencepiece" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/33/ea3cb3839607eb175da835244a798f797f478c5ddf0e8ecdf57ea85a4c70/sentencepiece-0.2.2.tar.gz", hash = "sha256:3d2b5e824b5622038dc7b490897efe05ebbbb9e7350fc142f3ecc8789ef9bdf6", size = 8218435, upload-time = "2026-07-12T08:39:34.701Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/1b/e6c69e4c2026ed575d68dda2847a404468ca7b5fa684bb0b19f71d82d29d/sentencepiece-0.2.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:bc7b0b1da20f856bfac5f84b2673fe534b167e41980b27442ca8f78c2b7eb77e", size = 2180607, upload-time = "2026-07-12T08:38:01.018Z" }, + { url = "https://files.pythonhosted.org/packages/1b/39/3d43a75dd5a22503ca5074d0d37707cabb2e4a71b4bc6e6c61be3643cc7a/sentencepiece-0.2.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8f1f61592e7cabd45d49ce8cc0ef42ca655c091e037153754fb3fa59725b5914", size = 1345667, upload-time = "2026-07-12T08:38:04.657Z" }, + { url = "https://files.pythonhosted.org/packages/90/d5/a69a8cc896e7de3fe2061b08c2f33e28656f243bed8af6a2df9f5d8c3124/sentencepiece-0.2.2-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c798f0b327bac10dc95cdac77b9a197ab2bd7dd1e60ebd7586a12d918d4be711", size = 1322864, upload-time = "2026-07-12T08:38:06.49Z" }, + { url = "https://files.pythonhosted.org/packages/e4/79/dd1836df32971d4eb14ff5cb4a8b3fe4419adbeada8e81d09dc53c5c0ef0/sentencepiece-0.2.2-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44284adc6fbe9d5bdd480541431a3d93f674fa44736714d3ad4bcee8283ace7d", size = 1392757, upload-time = "2026-07-12T08:38:08.559Z" }, + { url = "https://files.pythonhosted.org/packages/20/31/f23a2efaa0210b883574001b88fa64e499f798f0848a0b610fb9b384d162/sentencepiece-0.2.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:69e9dc8078e128286ed3b975e37c837ba96e215a50c3ef9f3f8b7ab9e5a832a0", size = 2184255, upload-time = "2026-07-12T08:38:14.855Z" }, + { url = "https://files.pythonhosted.org/packages/2a/92/3a6ea4a2c6dd9e7062698a5a33534ca0e20844883338ae9c6b9c122c1a9f/sentencepiece-0.2.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:443ac618c7a2a1377cf5c82581fbb849591d14e656d5e5a3e4682d4e36a34e4e", size = 1346997, upload-time = "2026-07-12T08:38:18.499Z" }, + { url = "https://files.pythonhosted.org/packages/f3/3a/7839048997c7bc0c34c57526f539f835e20c7a57dc2a99f99579b11cdbef/sentencepiece-0.2.2-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0e2aae42960392d6dcb9a72d8e1e65a97294c965071b43c7b3429a42f350250e", size = 1324282, upload-time = "2026-07-12T08:38:20.342Z" }, + { url = "https://files.pythonhosted.org/packages/06/5f/9117bf854aef817ad0d0ee9310eed0308a7e529e7eaf2e80ad9cd281ef82/sentencepiece-0.2.2-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1416b92f2f010333786fe6306ed2631121d5ea492219b0841e967b6765e64107", size = 1394242, upload-time = "2026-07-12T08:38:22.976Z" }, + { url = "https://files.pythonhosted.org/packages/b8/13/7a562289c8d5b49ebdf3f9c1e8ab67cf14a8743b1d90c8f406bfdec36b72/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1edb10e520e4bddf74d85b0f5ae74cc2d60c2b448885080bfb618bc2b3a49f6b", size = 2188384, upload-time = "2026-07-12T08:38:28.486Z" }, + { url = "https://files.pythonhosted.org/packages/bd/44/caa9cab5f261a019e2808bc5046152775dc57352ba9cbae7525e9e7a1ed4/sentencepiece-0.2.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38111ed1f79268f399c505028023d5eaaf0ab4e5eafceb709468b0d3323e7838", size = 1347176, upload-time = "2026-07-12T08:38:32.211Z" }, + { url = "https://files.pythonhosted.org/packages/19/90/cd798935668cff71d309d8ff10385844ecf216b1fe454f1993ed8bf2cb91/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cbce24284f51f71d10a42b7b9c964dcb9048b28f1c8e5db40bcbcb6f428cba6a", size = 1325200, upload-time = "2026-07-12T08:38:33.689Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2d/37e3da037318a70066ded0d51bc2a7f35491ae6338dd993d5eb1503fc3b5/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c8a168b040bc61681293f79a949b5d911c8e25086f4260285b8d97ab5f1195da", size = 1397736, upload-time = "2026-07-12T08:38:35.771Z" }, + { url = "https://files.pythonhosted.org/packages/b9/a3/b3b05095c174d6e80d37d5ddc2f57c2c56237333e7bbd6079cf3243c2a8a/sentencepiece-0.2.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:77c3ce990b23441e5ecfa5bce181fd6f408b564aeb6d7e1d1e7de9c5612501c8", size = 2188346, upload-time = "2026-07-12T08:38:41.089Z" }, + { url = "https://files.pythonhosted.org/packages/34/db/f9ea1a6844b4fa5dfe2312095cd866a1f724cd0905054ab9d5991778ba50/sentencepiece-0.2.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:201a8e0f55501a76e08dbf2c54bc45f4642b379271e89c667d517bfbc2191f2a", size = 1347267, upload-time = "2026-07-12T08:38:44.389Z" }, + { url = "https://files.pythonhosted.org/packages/32/4f/31c1073314ad94466bca37d29581761d70110237ee3d46b0efece59a8c1e/sentencepiece-0.2.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8eed98514bffe5ecac37f493f91869c351fbb05629328bfdbc08502c6c094dc0", size = 1324980, upload-time = "2026-07-12T08:38:46.304Z" }, + { url = "https://files.pythonhosted.org/packages/59/b4/a0356fa04d6a14337a6e0e443556785a0422c53ec58baae6b9568120eb0f/sentencepiece-0.2.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:64b656f025355cf8c51abe9fbe3848540756c6d7ca5e6791b1afa664bc24c7cb", size = 1397593, upload-time = "2026-07-12T08:38:48.302Z" }, + { url = "https://files.pythonhosted.org/packages/58/9c/dfc82846460e7a712310f5613f23d8b553cabb4e2e648663c11d8382af56/sentencepiece-0.2.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:72b7825b331b1b7e7c45be2e674b3e3c65af608fa376bad2d851b20aaf0cdc78", size = 2223080, upload-time = "2026-07-12T08:38:54.391Z" }, + { url = "https://files.pythonhosted.org/packages/59/5a/16d51d05360be4cee3ebfe4837c184054c4eed16cabaeb3b039524e9a000/sentencepiece-0.2.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3ab3f1ae98970b5590e2209341522718900ba19bcc2c207ffaa6bd417ad960c5", size = 1361138, upload-time = "2026-07-12T08:38:58.808Z" }, + { url = "https://files.pythonhosted.org/packages/0f/af/c30ee2a9f99d51db9844acaa8fa0b611a97c2fa7116646fa43db3300b187/sentencepiece-0.2.2-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ec27c152a1f1b24bc9168b55a5880f3c16e2334e697da6f55a1046a22405a3d", size = 1328625, upload-time = "2026-07-12T08:39:00.849Z" }, + { url = "https://files.pythonhosted.org/packages/3e/1a/4c6b39d03f5ba8439509adbd5a23c9538088a3cb679e7a47b911e8442bc6/sentencepiece-0.2.2-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:59d6588712101ccfcae9b03692be3aaae1514c2078666d7b05f15ba3a702e41b", size = 1398595, upload-time = "2026-07-12T08:39:02.86Z" }, +] + +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + +[[package]] +name = "tokenizers" +version = "0.23.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/18/1e/bc6587c5ab643b2e17776cace9070a2ae73549c86bffac9934a600bf3c31/tokenizers-0.23.2.tar.gz", hash = "sha256:7f0f085686b9de0d0079e6f874ae053600db64c5d13049e0bbc0119926d25aac", size = 385745, upload-time = "2026-09-03T08:55:42.89Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/49/22da045a91732384d3a3771816bf188dc5a1f702c32e635afa7c679c0bef/tokenizers-0.23.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:986670e43691469dcee610ea0f846f91a8f84e91fc6f7a48d4c064414c0ec2bf", size = 3101593, upload-time = "2026-09-03T08:55:28.587Z" }, + { url = "https://files.pythonhosted.org/packages/2e/4d/8f569ed49372a3ed8e57099bd515055fd48d7c95912c4307cda6973c2168/tokenizers-0.23.2-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a37039b5dfc4af84eb3ef0a92f4307e28936c8f9adccba2629d36f652e9bf7a2", size = 3516830, upload-time = "2026-09-03T08:55:14.741Z" }, + { url = "https://files.pythonhosted.org/packages/2a/de/e2f14c8919d5bf51874051d00d6c7b7e0e8bde6c6a2dbeddda7f642896ff/tokenizers-0.23.2-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7b7e37ba198f24150f523e1242e83c4970de4a525480586be5dcc24d9add32c5", size = 3407975, upload-time = "2026-09-03T08:55:16.842Z" }, + { url = "https://files.pythonhosted.org/packages/c5/bd/93c69152d02ef06ce47aed8b2bf4952dcf733c935a62791873932b2934d9/tokenizers-0.23.2-cp310-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:43e4f2071e3cc8d5d86421c874aebc82659bb51a68bcdef5a0da75ee89511ccb", size = 3748165, upload-time = "2026-09-03T08:55:24.769Z" }, + { url = "https://files.pythonhosted.org/packages/2d/b7/56b84b80bc96942bba8eb23751a9e8a1fce4faaf4390425e7083f721c98c/tokenizers-0.23.2-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:325fee2e0418a9dc6c9ecf736a5f5f0db7875183ace9549ae339da76f7a1fbb7", size = 4024165, upload-time = "2026-09-03T08:55:18.806Z" }, + { url = "https://files.pythonhosted.org/packages/9b/8a/0175e216f005c2fe08238292663aa41e4c802b216e71047a69a0e9fc6fa3/tokenizers-0.23.2-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:950d7c9426fa72406a0ffeacdbc0bb9985f5db20eb8b263f29c79aaf83105703", size = 3591899, upload-time = "2026-09-03T08:55:22.752Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ca/ca6b93c7820df123b2662a9469e8facc826ccc94e98fdd0d615f6431e73a/tokenizers-0.23.2-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:41c2f84d172449b4dadb9cdc508e3e364076613c35b16e76ecfe47a60d1e3305", size = 3386843, upload-time = "2026-09-03T08:55:26.584Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a4/4f9106d317b14a80aefea9f0e3a8d07ef25f856a7607eb7f5ab894281fcb/tokenizers-0.23.2-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:12f0835dc2ee694746a76adf7b1567d4346a4a502ebe93fb1f5f80ea49799b78", size = 3577314, upload-time = "2026-09-03T08:55:20.825Z" }, + { url = "https://files.pythonhosted.org/packages/8d/6a/1552b70fb0d9ab074fd3fc961435d01364e79c9058481822c3af6e8d402c/tokenizers-0.23.2-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:eb2f9c8a24da020ea8c11a01a19c1c2547912d92121ae4a01cfbca46125dee40", size = 9967367, upload-time = "2026-09-03T08:55:33.188Z" }, + { url = "https://files.pythonhosted.org/packages/06/01/3ccb3a956c7528b2507b8a9714155c4baf86af593039db6ea375dd0c96c3/tokenizers-0.23.2-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:f486f402f6f9abee5bb032553736813af0c710a86b2e0ca592634c55cea1f835", size = 9811886, upload-time = "2026-09-03T08:55:35.642Z" }, + { url = "https://files.pythonhosted.org/packages/fa/73/7038e612d48bda1599457f712f6bd3854eae1a9dc9c13aa47f835349db48/tokenizers-0.23.2-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:bef235815a067b2648caf6dcc7a71091b0b0fff9ee8057f6451eb9335fae52ef", size = 10146224, upload-time = "2026-09-03T08:55:38.391Z" }, + { url = "https://files.pythonhosted.org/packages/b5/d8/8e9e4e0b287a338d8f88976729628c9d22e8a54cfaf9777018a7f7cb58a0/tokenizers-0.23.2-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5c56bda1511921587789163e524d196ed8284174ac23abd7685d5ea8da6c4718", size = 10256304, upload-time = "2026-09-03T08:55:40.977Z" }, +] + +[[package]] +name = "tomli" +version = "2.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/22/de/48c59722572767841493b26183a0d1cc411d54fd759c5607c4590b6563a6/tomli-2.4.1.tar.gz", hash = "sha256:7c7e1a961a0b2f2472c1ac5b69affa0ae1132c39adcb67aba98568702b9cc23f", size = 17543, upload-time = "2026-03-25T20:22:03.828Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/f7/675db52c7e46064a9aa928885a9b20f4124ecb9bc2e1ce74c9106648d202/tomli-2.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4ab97e64ccda8756376892c53a72bd1f964e519c77236368527f758fbc36a53a", size = 149454, upload-time = "2026-03-25T20:21:12.036Z" }, + { url = "https://files.pythonhosted.org/packages/61/71/81c50943cf953efa35bce7646caab3cf457a7d8c030b27cfb40d7235f9ee/tomli-2.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96481a5786729fd470164b47cdb3e0e58062a496f455ee41b4403be77cb5a076", size = 237561, upload-time = "2026-03-25T20:21:13.098Z" }, + { url = "https://files.pythonhosted.org/packages/48/c1/f41d9cb618acccca7df82aaf682f9b49013c9397212cb9f53219e3abac37/tomli-2.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5a881ab208c0baf688221f8cecc5401bd291d67e38a1ac884d6736cbcd8247e9", size = 243824, upload-time = "2026-03-25T20:21:14.569Z" }, + { url = "https://files.pythonhosted.org/packages/22/e4/5a816ecdd1f8ca51fb756ef684b90f2780afc52fc67f987e3c61d800a46d/tomli-2.4.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:47149d5bd38761ac8be13a84864bf0b7b70bc051806bc3669ab1cbc56216b23c", size = 242227, upload-time = "2026-03-25T20:21:15.712Z" }, + { url = "https://files.pythonhosted.org/packages/6b/49/2b2a0ef529aa6eec245d25f0c703e020a73955ad7edf73e7f54ddc608aa5/tomli-2.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:ec9bfaf3ad2df51ace80688143a6a4ebc09a248f6ff781a9945e51937008fcbc", size = 247859, upload-time = "2026-03-25T20:21:17.001Z" }, + { url = "https://files.pythonhosted.org/packages/dc/c7/62d7a17c26487ade21c5422b646110f2162f1fcc95980ef7f63e73c68f14/tomli-2.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7f86fd587c4ed9dd76f318225e7d9b29cfc5a9d43de44e5754db8d1128487085", size = 150018, upload-time = "2026-03-25T20:21:23.002Z" }, + { url = "https://files.pythonhosted.org/packages/5c/05/79d13d7c15f13bdef410bdd49a6485b1c37d28968314eabee452c22a7fda/tomli-2.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ff18e6a727ee0ab0388507b89d1bc6a22b138d1e2fa56d1ad494586d61d2eae9", size = 244948, upload-time = "2026-03-25T20:21:24.04Z" }, + { url = "https://files.pythonhosted.org/packages/10/90/d62ce007a1c80d0b2c93e02cab211224756240884751b94ca72df8a875ca/tomli-2.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:136443dbd7e1dee43c68ac2694fde36b2849865fa258d39bf822c10e8068eac5", size = 253341, upload-time = "2026-03-25T20:21:25.177Z" }, + { url = "https://files.pythonhosted.org/packages/1a/7e/caf6496d60152ad4ed09282c1885cca4eea150bfd007da84aea07bcc0a3e/tomli-2.4.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5e262d41726bc187e69af7825504c933b6794dc3fbd5945e41a79bb14c31f585", size = 248159, upload-time = "2026-03-25T20:21:26.364Z" }, + { url = "https://files.pythonhosted.org/packages/99/e7/c6f69c3120de34bbd882c6fba7975f3d7a746e9218e56ab46a1bc4b42552/tomli-2.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5cb41aa38891e073ee49d55fbc7839cfdb2bc0e600add13874d048c94aadddd1", size = 253290, upload-time = "2026-03-25T20:21:27.46Z" }, + { url = "https://files.pythonhosted.org/packages/14/6f/12645cf7f08e1a20c7eb8c297c6f11d31c1b50f316a7e7e1e1de6e2e7b7e/tomli-2.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:eb0dc4e38e6a1fd579e5d50369aa2e10acfc9cace504579b2faabb478e76941a", size = 149887, upload-time = "2026-03-25T20:21:33.028Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e0/90637574e5e7212c09099c67ad349b04ec4d6020324539297b634a0192b0/tomli-2.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7f2c7f2b9ca6bdeef8f0fa897f8e05085923eb091721675170254cbc5b02897", size = 243704, upload-time = "2026-03-25T20:21:34.51Z" }, + { url = "https://files.pythonhosted.org/packages/10/8f/d3ddb16c5a4befdf31a23307f72828686ab2096f068eaf56631e136c1fdd/tomli-2.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3c6818a1a86dd6dca7ddcaaf76947d5ba31aecc28cb1b67009a5877c9a64f3f", size = 251628, upload-time = "2026-03-25T20:21:36.012Z" }, + { url = "https://files.pythonhosted.org/packages/e3/f1/dbeeb9116715abee2485bf0a12d07a8f31af94d71608c171c45f64c0469d/tomli-2.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d312ef37c91508b0ab2cee7da26ec0b3ed2f03ce12bd87a588d771ae15dcf82d", size = 247180, upload-time = "2026-03-25T20:21:37.136Z" }, + { url = "https://files.pythonhosted.org/packages/d3/74/16336ffd19ed4da28a70959f92f506233bd7cfc2332b20bdb01591e8b1d1/tomli-2.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:51529d40e3ca50046d7606fa99ce3956a617f9b36380da3b7f0dd3dd28e68cb5", size = 251674, upload-time = "2026-03-25T20:21:38.298Z" }, + { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, +] + +[[package]] +name = "tqdm" +version = "4.70.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/ea/b2a5bd54b28a324dae8211928b2d730b6547500342c7e6c6dea08bd0a485/tqdm-4.70.1.tar.gz", hash = "sha256:cefd0eca11b2a37a3aee776544d4f4ae913f02688135b5556b8788dfa474afc4", size = 171846, upload-time = "2026-09-11T07:25:16.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/03/921a3d3c75785aca9ebfbfcabfbc3a1be12e2ab5265deb026d55a5a3f83e/tqdm-4.70.1-py3-none-any.whl", hash = "sha256:c293e525e6fef9c20e8728fd4612df02a0aa31bb5fe91ecd93e123b1b7bffa73", size = 80199, upload-time = "2026-09-11T07:25:14.599Z" }, +] + +[[package]] +name = "transformers" +version = "5.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "numpy", version = "2.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/9e/750649904a065007a838981785b2bd8d9ff26154c6c341ac67d0b7f82c68/transformers-5.17.0.tar.gz", hash = "sha256:a153be279169b55b92d8000bf4af294aed684503d091cca7804da2dd8a9de000", size = 9817878, upload-time = "2026-09-09T15:39:56.886Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/d0/c502b60d684adbd98a8dc7d5bb866842772b816ac4354e4608be240041ae/transformers-5.17.0-py3-none-any.whl", hash = "sha256:78ec1ce21579b38dfb83950a0658cd119f87212a2fcfdff478096ce9d6c03801", size = 12295140, upload-time = "2026-09-09T15:39:53.746Z" }, +] + +[[package]] +name = "trida-mlx" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "huggingface-hub" }, + { name = "mlx" }, + { name = "mlx", extra = ["cpu"], marker = "sys_platform == 'linux'" }, + { name = "mlx-lm" }, + { name = "transformers" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "huggingface-hub" }, + { name = "mlx", marker = "sys_platform == 'darwin'", specifier = ">=0.29" }, + { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = ">=0.29" }, + { name = "mlx-lm", specifier = ">=0.31" }, + { name = "transformers", specifier = ">=4.57" }, +] + +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=8" }] + +[[package]] +name = "typer" +version = "0.27.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "rich" }, + { name = "shellingham" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/16/f7/57713ba479fd405eb76de31404b2c744c289e336b2d999511ebf51e496f7/typer-0.27.2.tar.gz", hash = "sha256:269b7eb9d3c202ca84b4bc9618cb04ebb43d3d4d1e567e4c768607232c05f945", size = 204045, upload-time = "2026-08-28T10:26:55.046Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/bf/205d0004930ede8f542fb58f601526fccf4ae7626075ca1e6c4de5d3d652/typer-0.27.2-py3-none-any.whl", hash = "sha256:b3a5fc4342d5fc8fda8fc3010b1cf117e9249aab7fae800c2eff62fd3842d97d", size = 123130, upload-time = "2026-08-28T10:26:53.752Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] From 285da1fa1f690ee9553885088086ad30efc9eaf4 Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Wed, 30 Sep 2026 01:29:25 +0000 Subject: [PATCH 2/9] feat(inference/mlx): OpenAI-compatible server, agent loop, Hermes support - /v1/chat/completions and /v1/completions with SSE streaming, reasoning split, qwen3-coder / JSON tool calls parsed into OpenAI tool_calls with schema-typed parameters, stop sequences, per-request thinking control - /v1/models advertises context_length; context_length_exceeded 400s - all MLX work on one worker thread (MLX streams are per-thread; the threaded HTTP server crashed on chunked prefills without it) - Hermes Agent compatibility: template-safe message normalisation, prefill keep-alives, 3 cache slots + a system-block snapshot so side requests and rewritten user turns keep the tool-schema prefix cached - minimal stdlib agent loop (files, search, calculator, opt-in shell) - offline end-to-end tests on a tiny raw-HF checkpoint (loader, convert q8, prompt cache, server, parsing, the real Trida chat template) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- .../tests/fixtures/trida_chat_template.jinja | 154 ++++++ inference/mlx/tests/test_e2e.py | 229 +++++++++ inference/mlx/tests/tiny_ckpt.py | 96 ++++ inference/mlx/trida_mlx/agent.py | 222 +++++++++ inference/mlx/trida_mlx/parsing.py | 222 +++++++++ inference/mlx/trida_mlx/server.py | 462 ++++++++++++++++++ 6 files changed, 1385 insertions(+) create mode 100644 inference/mlx/tests/fixtures/trida_chat_template.jinja create mode 100644 inference/mlx/tests/test_e2e.py create mode 100644 inference/mlx/tests/tiny_ckpt.py create mode 100644 inference/mlx/trida_mlx/agent.py create mode 100644 inference/mlx/trida_mlx/parsing.py create mode 100644 inference/mlx/trida_mlx/server.py diff --git a/inference/mlx/tests/fixtures/trida_chat_template.jinja b/inference/mlx/tests/fixtures/trida_chat_template.jinja new file mode 100644 index 0000000..ade1daa --- /dev/null +++ b/inference/mlx/tests/fixtures/trida_chat_template.jinja @@ -0,0 +1,154 @@ +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {%- if content %} + {{- '\n\n' + content }} + {%- endif %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if ns.multi_step_tool %} + {{- raise_exception('No user query found in messages.') }} +{%- endif %} +{%- for message in messages %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "system" %} + {%- if not loop.first %} + {{- raise_exception('System message must be at the beginning.') }} + {%- endif %} + {%- elif message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {%- if (loop.index0 > ns.last_query_index) or (keep_all_reasoning is defined and keep_all_reasoning) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- else %} + {{- raise_exception('Unexpected message role.') }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} diff --git a/inference/mlx/tests/test_e2e.py b/inference/mlx/tests/test_e2e.py new file mode 100644 index 0000000..d054fbc --- /dev/null +++ b/inference/mlx/tests/test_e2e.py @@ -0,0 +1,229 @@ +"""End-to-end (offline) checks on a tiny raw-HF checkpoint: loader, convert (q8), engine +prompt cache, OpenAI server (non-stream + SSE + tools), and output parsing.""" +import json +import os +import socket +import sys +import tempfile +import threading +import urllib.error +import urllib.request +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import mlx.core as mx + +import tiny_ckpt +from test_tiny import tiny_runtime +from trida_mlx import convert +from trida_mlx.decode import DecodeStats, SamplingParams +from trida_mlx.engine import Engine +from trida_mlx.model import TridaRuntime, load_model +from trida_mlx.parsing import StreamSplitter, parse_tool_calls, split_reasoning + +TMP = Path(tempfile.mkdtemp(prefix="trida_e2e_")) +CKPT, MASK_ID = tiny_ckpt.build(TMP / "raw") + + +def test_raw_hf_loader_matches_reference(): + model, cfg, _ = load_model(str(CKPT), dtype=mx.float32) + rt = TridaRuntime(model) + ref = tiny_runtime(seed=0) + p = [5, 17, 99, 3, 250, 8, 41] + a = rt.prefill(rt.make_cache(), p) + b = ref.prefill(ref.make_cache(), p) + # bf16 round trip of the weights -> small differences only + assert float(mx.max(mx.abs(a - b))) < 0.25 * float(mx.max(mx.abs(b))), (a[:5], b[:5]) + assert int(mx.argmax(a)) == int(mx.argmax(b)) + + +def test_convert_q8_and_selfspec_lossless_on_quantized(): + out = TMP / "q8" + convert.main(["--model", str(CKPT), "--out", str(out), "--bits", "8", "--group-size", "32"]) + eng = Engine(str(out)) + assert eng.mask_id == MASK_ID + p = eng.encode(eng.render([{"role": "user", "content": "hello world"}])) + sp = SamplingParams(temperature=0.0) + a = [t for c in eng.generate_ids(p, max_tokens=80, sampling=sp, mode="causal") for t in c] + eng.reset_cache() + b = [t for c in eng.generate_ids(p, max_tokens=80, sampling=sp, mode="self-spec") for t in c] + assert a == b + + +def test_prompt_cache_reuse(): + eng = Engine(str(CKPT)) + sp = SamplingParams(temperature=0.0) + msgs = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello world"}] + p1 = eng.encode(eng.render(msgs)) + out1 = [t for c in eng.generate_ids(p1, max_tokens=30, sampling=sp) for t in c] + msgs2 = msgs + [{"role": "assistant", "content": eng.tokenizer.decode(out1)}, + {"role": "user", "content": "and the weather?"}] + p2 = eng.encode(eng.render(msgs2)) + st = DecodeStats() + out2 = [t for c in eng.generate_ids(p2, max_tokens=30, sampling=sp, stats=st) for t in c] + assert st.reused_tokens > len(p1) // 2, st.as_dict() + # same result as a cold engine + eng2 = Engine(str(CKPT), prompt_cache=False) + ref2 = [t for c in eng2.generate_ids(p2, max_tokens=30, sampling=sp) for t in c] + assert out2 == ref2 + + +def test_parsing(): + c, calls = parse_tool_calls('ok {"name": "f", "arguments": {"x": 1}}}') + assert c == "ok" and calls == [{"name": "f", "arguments": {"x": 1}}] + _, calls = parse_tool_calls("\n\n\nSeoul\n\n\n") + assert calls == [{"name": "g", "arguments": {"city": "Seoul"}}] + assert split_reasoning("abc\n\nhi", True) == ("abc", "hi") + s = StreamSplitter(True, parse_tools=True) + text = 'let me think\n\nSure.{"name": "f", "arguments": {}}' + parts = [] + for ch in text: + parts += s.feed(ch) + rest, calls = s.finish() + parts += rest + reasoning = "".join(t for k, t in parts if k == "reasoning") + content = "".join(t for k, t in parts if k == "content") + assert reasoning == "let me think" and content == "Sure." and calls[0]["name"] == "f", (parts, calls) + + +def test_normalize_messages_against_real_trida_template(): + """Hermes-style histories must render with the real Trida chat template.""" + from transformers import AutoTokenizer + from trida_mlx.server import _normalize_messages + tok = AutoTokenizer.from_pretrained(str(CKPT)) + tok.chat_template = (Path(__file__).parent / "fixtures" / "trida_chat_template.jinja").read_text() + tools = [{"type": "function", "function": {"name": "terminal", "description": "run", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}}}] + msgs = [ + {"role": "system", "content": "You are Hermes."}, + {"role": "developer", "content": "Be terse."}, + {"role": "user", "content": [{"type": "text", "text": "list files"}]}, + {"role": "assistant", "content": None, "reasoning": "use terminal", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "terminal", "arguments": '{"command": "ls"}'}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "a.py\nb.py"}, + {"role": "system", "content": "Context was compacted."}, + {"role": "user", "content": "thanks, now count them"}, + ] + norm = _normalize_messages(msgs) + assert [m["role"] for m in norm] == ["system", "user", "assistant", "tool", "user", "user"] + text = tok.apply_chat_template(norm, tools=tools, add_generation_prompt=True, tokenize=False, enable_thinking=True) + assert "You are Hermes.\n\nBe terse." in text + assert "\n\nls\n" in text + assert "[System note]\nContext was compacted." in text and text.endswith("\n") + + +def test_tool_param_coercion(): + tools = [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": { + "path": {"type": "string"}, "n": {"type": "integer"}, "flags": {"type": "array"}}}}}] + _, calls = parse_tool_calls("\n\n\n007\n\n" + "\n5\n\n\n[\"-a\"]\n\n" + "\n", tools) + assert calls == [{"name": "f", "arguments": {"path": "007", "n": 5, "flags": ["-a"]}}], calls + + +def test_multi_slot_cache_survives_side_requests(): + eng = Engine(str(CKPT), cache_slots=3) + sp = SamplingParams(temperature=0.0) + run = lambda msgs, st=None: [t for c in eng.generate_ids(eng.encode(eng.render(msgs)), max_tokens=12, + sampling=sp, stats=st) for t in c] + main = [{"role": "system", "content": "main agent " * 30}, {"role": "user", "content": "hello world"}] + out = run(main) + run([{"role": "user", "content": "write a title for: hello"}]) # side request (title) + run([{"role": "user", "content": "summarize: the weather in seoul"}]) # another side request + st = DecodeStats() + run(main + [{"role": "assistant", "content": eng.tokenizer.decode(out)}, + {"role": "user", "content": "and then?"}], st) + assert st.reused_tokens > 100, st.as_dict() + assert len(eng.slots) == 3 + + +def test_system_block_snapshot_reused_when_user_turn_changes(): + """Hermes rewrites the last user message on a truncation retry; the (long) system block + must still be reused.""" + eng = Engine(str(CKPT), cache_slots=1) + sp = SamplingParams(temperature=0.0) + sys_msg = {"role": "system", "content": "tool schemas " * 60} + p1 = eng.encode(eng.render([sys_msg, {"role": "user", "content": "list the files"}])) + for _ in eng.generate_ids(p1, max_tokens=5, sampling=sp): + pass + st = DecodeStats() + p2 = eng.encode(eng.render([sys_msg, {"role": "user", "content": "list the files [System: continue]"}])) + out = [t for c in eng.generate_ids(p2, max_tokens=5, sampling=sp, stats=st) for t in c] + n_sys = p1.index(eng.im_start, 1) # tokens up to the end of the system block + assert st.reused_tokens == n_sys, (st.reused_tokens, n_sys) + ref = Engine(str(CKPT), prompt_cache=False) + assert out == [t for c in ref.generate_ids(p2, max_tokens=5, sampling=sp) for t in c] + + +def _free_port(): + s = socket.socket(); s.bind(("127.0.0.1", 0)); p = s.getsockname()[1]; s.close(); return p + + +def test_server_roundtrip(): + from trida_mlx import server + port = _free_port() + t = threading.Thread(target=server.main, args=(["--model", str(CKPT), "--port", str(port), + "--max-tokens", "40", "--prefill-chunk", "16", + "--context-length", "2048"],), daemon=True) + t.start() + base = f"http://127.0.0.1:{port}/v1" + import time + for _ in range(200): + try: + urllib.request.urlopen(f"{base}/models", timeout=1); break + except Exception: + time.sleep(0.1) + tools = [{"type": "function", "function": {"name": "now", "description": "time", + "parameters": {"type": "object", "properties": {}}}}] + body = {"messages": [{"role": "user", "content": "hello"}], "tools": tools, "temperature": 0} + r = json.loads(urllib.request.urlopen(urllib.request.Request( + f"{base}/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"})).read()) + assert r["choices"][0]["message"]["role"] == "assistant" + assert r["usage"]["trida"]["new_tokens"] > 0 + body["stream"] = True + lines = urllib.request.urlopen(urllib.request.Request( + f"{base}/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"})).read().decode().splitlines() + datas = [l[6:] for l in lines if l.startswith("data: ")] + assert datas[-1] == "[DONE]" + usage = json.loads(datas[-2])["usage"] + assert usage["prompt_tokens_details"]["cached_tokens"] > 0 # second identical request hits the cache + assert any(l.startswith(": prefill") for l in lines) or usage["prompt_tokens_details"]["cached_tokens"] > 0 + r = json.loads(urllib.request.urlopen(urllib.request.Request( + f"{base}/completions", data=json.dumps({"prompt": "hello", "max_tokens": 10}).encode(), + headers={"Content-Type": "application/json"})).read()) + assert r["usage"]["completion_tokens"] <= 10 + models = json.loads(urllib.request.urlopen(f"{base}/models").read()) + assert models["data"][0]["context_length"] == 2048 and models["data"][0]["id"] == "trida2.0-4b" + # context overflow -> OpenAI-style 400 that harnesses recognise (and compress) + big = {"messages": [{"role": "user", "content": "hello world " * 3000}], "stream": True} + try: + urllib.request.urlopen(urllib.request.Request(f"{base}/chat/completions", data=json.dumps(big).encode(), + headers={"Content-Type": "application/json"})) + assert False, "expected 400" + except urllib.error.HTTPError as e: + err = json.loads(e.read())["error"] + assert e.code == 400 and err["code"] == "context_length_exceeded" and "maximum context length" in err["message"] + # keep-alive comments during a long (uncached) prefill + body2 = {"messages": [{"role": "user", "content": "the weather in seoul is sunny " * 40}], "stream": True, + "max_tokens": 4} + lines = urllib.request.urlopen(urllib.request.Request( + f"{base}/chat/completions", data=json.dumps(body2).encode(), + headers={"Content-Type": "application/json"})).read().decode().splitlines() + assert any(l.startswith(": prefill") for l in lines), lines[:5] + + +if __name__ == "__main__": + fails = 0 + for name, fn in list(globals().items()): + if name.startswith("test_"): + try: + fn(); print(f" ok {name}") + except Exception as e: # noqa: BLE001 + import traceback; traceback.print_exc() + fails += 1; print(f" FAIL {name}: {type(e).__name__}: {e}") + sys.stdout.flush() + os._exit(1 if fails else 0) diff --git a/inference/mlx/tests/tiny_ckpt.py b/inference/mlx/tests/tiny_ckpt.py new file mode 100644 index 0000000..31474c0 --- /dev/null +++ b/inference/mlx/tests/tiny_ckpt.py @@ -0,0 +1,96 @@ +"""Build a tiny *raw-HF-format* Trida-like checkpoint (random weights, byte-level BPE +tokenizer, Qwen-style chat template with tools + thinking) for offline end-to-end tests.""" +import json +import os +import sys +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import mlx.core as mx +from mlx.utils import tree_flatten + +from test_tiny import TINY, tiny_runtime # noqa: E402 + +SPECIALS = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "", "", "", + "", "<|mask|>"] + +CHAT_TEMPLATE = r"""{%- if tools %}<|im_start|>system +{%- if messages[0].role == 'system' %} +{{ messages[0].content }} +{%- endif %} +# Tools + +{%- for t in tools %} +{{ t | tojson }} +{%- endfor %} + +Call a tool with {"name": ..., "arguments": {...}}<|im_end|> +{% elif messages[0].role == 'system' %}<|im_start|>system +{{ messages[0].content }}<|im_end|> +{% endif %} +{%- for m in messages %} +{%- if m.role == 'user' %}<|im_start|>user +{{ m.content }}<|im_end|> +{% elif m.role == 'assistant' %}<|im_start|>assistant +{{ m.content }} +{%- if m.tool_calls %}{% for tc in m.tool_calls %}{"name": "{{ tc.function.name }}", "arguments": {{ tc.function.arguments | tojson }}}{% endfor %}{% endif %}<|im_end|> +{% elif m.role == 'tool' %}<|im_start|>user +{{ m.content }}<|im_end|> +{% endif %} +{%- endfor %} +{%- if add_generation_prompt %}<|im_start|>assistant +{%- if enable_thinking is defined and enable_thinking is false %} + + + + +{% else %} + +{% endif %} +{%- endif %}""" + + +def build(out: Path): + from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers + from transformers import PreTrainedTokenizerFast + + out.mkdir(parents=True, exist_ok=True) + tok = Tokenizer(models.BPE()) + tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + tok.decoder = decoders.ByteLevel() + corpus = ["hello world the weather in seoul is sunny", "def fib(n): return n", "tool call json name", + "안녕하세요 온디바이스"] * 20 + trainer = trainers.BpeTrainer(vocab_size=TINY["vocab_size"] - len(SPECIALS), special_tokens=[], + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False) + tok.train_from_iterator(corpus, trainer) + tok.add_special_tokens(SPECIALS) + assert tok.get_vocab_size() <= TINY["vocab_size"], tok.get_vocab_size() + fast = PreTrainedTokenizerFast(tokenizer_object=tok, eos_token="<|im_end|>", pad_token="<|endoftext|>") + fast.chat_template = CHAT_TEMPLATE + fast.save_pretrained(str(out)) + mask_id = fast.convert_tokens_to_ids("<|mask|>") + + rt = tiny_runtime(seed=0) + params = dict(tree_flatten(rt.model.parameters())) + hf = {} + norm_sfx = (".input_layernorm.weight", ".post_attention_layernorm.weight", "model.norm.weight", + ".q_norm.weight", ".k_norm.weight") + for k, v in params.items(): + k2 = k.replace("language_model.model.", "model.").replace("language_model.", "") + if "conv1d.weight" in k2: + v = v.moveaxis(1, 2) # MLX (C, K, 1) -> HF (C, 1, K) + if any(k2.endswith(s) for s in norm_sfx): + v = v - 1.0 # HF Qwen3.5 stores zero-centred RMSNorm weights + hf[k2] = v.astype(mx.bfloat16) if not k2.endswith("A_log") else v + hf["lm_head.weight"] = hf["model.embed_tokens.weight"] # real ckpt ships the tied head too + mx.save_safetensors(str(out / "model.safetensors"), hf) + cfg = dict(TINY, architectures=["Qwen3_5ForCausalLM"]) + (out / "config.json").write_text(json.dumps(cfg, indent=2)) + (out / "block_diffusion.json").write_text(json.dumps({"bd_size": 4, "mask_id": mask_id})) + (out / "generation_config.json").write_text(json.dumps({"eos_token_id": [fast.eos_token_id]})) + return out, mask_id + + +if __name__ == "__main__": + print(build(Path(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tiny_trida"))) diff --git a/inference/mlx/trida_mlx/agent.py b/inference/mlx/trida_mlx/agent.py new file mode 100644 index 0000000..50830d3 --- /dev/null +++ b/inference/mlx/trida_mlx/agent.py @@ -0,0 +1,222 @@ +"""A minimal local agent loop against the on-device server (stdlib only). + + python -m trida_mlx.server --model ./Trida2.0-4B-mlx-q8 & # terminal 1 + python -m trida_mlx.agent --workdir ~/some/project # terminal 2 + python -m trida_mlx.agent --workdir . --allow-shell "how many python files are here?" + +Tools: list_dir, read_file, search_files, calculator, now, and (opt-in) run_shell / +write_file — both ask for confirmation unless --yes. Any OpenAI-compatible client works +the same way; this file only exists to exercise tool calling end to end on device. +""" +from __future__ import annotations + +import argparse +import ast +import datetime as _dt +import json +import operator +import os +import subprocess +import sys +import urllib.request +from pathlib import Path + +DIM, BOLD, CYAN, RESET = "\033[2m", "\033[1m", "\033[36m", "\033[0m" + + +def _fn(name, desc, props, required=()): + return {"type": "function", "function": {"name": name, "description": desc, "parameters": { + "type": "object", "properties": props, "required": list(required)}}} + + +class Tools: + def __init__(self, workdir: Path, allow_shell: bool, allow_write: bool, yes: bool): + self.root, self.allow_shell, self.allow_write, self.yes = workdir.resolve(), allow_shell, allow_write, yes + + def specs(self): + s = [ + _fn("list_dir", "List files in a directory (relative to the workspace).", {"path": {"type": "string"}}), + _fn("read_file", "Read a text file (relative to the workspace). Returns at most max_chars.", + {"path": {"type": "string"}, "max_chars": {"type": "integer"}}, ["path"]), + _fn("search_files", "Find lines matching a substring in files under a directory.", + {"query": {"type": "string"}, "path": {"type": "string"}}, ["query"]), + _fn("calculator", "Evaluate an arithmetic expression, e.g. '17*24+3'.", + {"expression": {"type": "string"}}, ["expression"]), + _fn("now", "Current local date and time.", {}), + ] + if self.allow_shell: + s.append(_fn("run_shell", "Run a shell command in the workspace; returns stdout/stderr.", + {"command": {"type": "string"}}, ["command"])) + if self.allow_write: + s.append(_fn("write_file", "Write a text file (relative to the workspace).", + {"path": {"type": "string"}, "content": {"type": "string"}}, ["path", "content"])) + return s + + def _p(self, rel): + p = (self.root / (rel or ".")).resolve() + if self.root not in (p, *p.parents): + raise ValueError("path escapes the workspace") + return p + + def _confirm(self, what): + if self.yes: + return True + return input(f"{BOLD}allow {what}? [y/N] {RESET}").strip().lower() == "y" + + def call(self, name, args): + try: + if name == "list_dir": + p = self._p(args.get("path")) + return "\n".join(sorted((e.name + ("/" if e.is_dir() else "")) for e in p.iterdir())[:500]) + if name == "read_file": + return self._p(args["path"]).read_text(errors="replace")[: int(args.get("max_chars") or 8000)] + if name == "search_files": + q, hits = args["query"], [] + for f in self._p(args.get("path")).rglob("*"): + if f.is_file() and f.stat().st_size < 1_000_000 and ".git" not in f.parts: + try: + for i, line in enumerate(f.read_text(errors="ignore").splitlines(), 1): + if q in line: + hits.append(f"{f.relative_to(self.root)}:{i}: {line.strip()[:200]}") + except Exception: + pass + if len(hits) >= 50: + break + return "\n".join(hits) or "(no matches)" + if name == "calculator": + return str(_safe_eval(args["expression"])) + if name == "now": + return _dt.datetime.now().isoformat(timespec="seconds") + if name == "run_shell" and self.allow_shell: + if not self._confirm(f"shell `{args['command']}`"): + return "user denied" + r = subprocess.run(args["command"], shell=True, cwd=self.root, capture_output=True, text=True, + timeout=60) + return (r.stdout + ("\n[stderr]\n" + r.stderr if r.stderr else ""))[-8000:] or f"(exit {r.returncode})" + if name == "write_file" and self.allow_write: + p = self._p(args["path"]) + if not self._confirm(f"write {p}"): + return "user denied" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(args["content"]) + return f"wrote {len(args['content'])} chars to {args['path']}" + return f"unknown tool {name}" + except Exception as e: # noqa: BLE001 + return f"error: {type(e).__name__}: {e}" + + +_OPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, + ast.Pow: operator.pow, ast.Mod: operator.mod, ast.FloorDiv: operator.floordiv, ast.USub: operator.neg} + + +def _safe_eval(expr): + def ev(n): + if isinstance(n, ast.Constant) and isinstance(n.value, (int, float)): + return n.value + if isinstance(n, ast.BinOp) and type(n.op) in _OPS: + return _OPS[type(n.op)](ev(n.left), ev(n.right)) + if isinstance(n, ast.UnaryOp) and type(n.op) in _OPS: + return _OPS[type(n.op)](ev(n.operand)) + raise ValueError("unsupported expression") + return ev(ast.parse(expr, mode="eval").body) + + +def stream_chat(base, body): + req = urllib.request.Request(f"{base}/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + msg = {"role": "assistant", "content": "", "reasoning_content": "", "tool_calls": []} + usage, in_reason = None, False + with urllib.request.urlopen(req) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + data = line[6:] + if data == "[DONE]": + break + obj = json.loads(data) + if obj.get("usage"): + usage = obj["usage"] + for ch in obj.get("choices", []): + d = ch.get("delta", {}) + if d.get("reasoning_content"): + if not in_reason: + sys.stdout.write(DIM); in_reason = True + sys.stdout.write(d["reasoning_content"]); msg["reasoning_content"] += d["reasoning_content"] + if d.get("content"): + if in_reason: + sys.stdout.write(RESET + "\n"); in_reason = False + sys.stdout.write(d["content"]); msg["content"] += d["content"] + if d.get("tool_calls"): + msg["tool_calls"].extend(d["tool_calls"]) + sys.stdout.flush() + if in_reason: + sys.stdout.write(RESET) + print() + return msg, usage + + +def main(argv=None): + ap = argparse.ArgumentParser() + ap.add_argument("task", nargs="*", help="one-shot task; omit for an interactive session") + ap.add_argument("--base-url", default="http://127.0.0.1:8080/v1") + ap.add_argument("--workdir", default=".") + ap.add_argument("--allow-shell", action="store_true") + ap.add_argument("--allow-write", action="store_true") + ap.add_argument("--yes", action="store_true", help="don't ask before shell/write tools") + ap.add_argument("--max-steps", type=int, default=12) + ap.add_argument("--no-think", action="store_true") + ap.add_argument("--temperature", type=float, default=None) + a = ap.parse_args(argv) + tools = Tools(Path(a.workdir).expanduser(), a.allow_shell, a.allow_write, a.yes) + messages = [{"role": "system", "content": + "You are a helpful on-device assistant. Use the provided tools when they help; the workspace " + f"root is {tools.root}. Paths are relative to it. Be concise."}] + + def run(task): + messages.append({"role": "user", "content": task}) + for step in range(a.max_steps): + body = {"messages": messages, "tools": tools.specs(), "stream": True, + "chat_template_kwargs": {"enable_thinking": not a.no_think}} + if a.temperature is not None: + body["temperature"] = a.temperature + msg, usage = stream_chat(a.base_url, body) + if usage and usage.get("trida"): + t = usage["trida"] + print(f"{DIM}[{t['mode']}: {t['new_tokens']} tok, {t['decode_tok_s']} tok/s, " + f"{t['tokens_per_forward']} tok/fwd, prompt {t['prompt_tokens']} " + f"(cached {t['reused_tokens']}), prefill {t['prefill_s']}s]{RESET}") + entry = {"role": "assistant", "content": msg["content"]} + if msg["tool_calls"]: + entry["tool_calls"] = [{"id": tc["id"], "type": "function", "function": tc["function"]} + for tc in msg["tool_calls"]] + messages.append(entry) + if not msg["tool_calls"]: + return + for tc in msg["tool_calls"]: + fn = tc["function"] + try: + args = json.loads(fn.get("arguments") or "{}") + except Exception: + args = {} + print(f"{CYAN}→ {fn['name']}({json.dumps(args, ensure_ascii=False)[:200]}){RESET}") + result = tools.call(fn["name"], args) + print(f"{DIM}{result[:600]}{'…' if len(result) > 600 else ''}{RESET}") + messages.append({"role": "tool", "tool_call_id": tc["id"], "name": fn["name"], "content": result}) + print("(max steps reached)") + + if a.task: + run(" ".join(a.task)) + return + print(f"{BOLD}trida on-device agent{RESET} — workspace {tools.root}. Ctrl-D to quit.") + while True: + try: + task = input(f"{BOLD}> {RESET}").strip() + except EOFError: + break + if task: + run(task) + + +if __name__ == "__main__": + main() diff --git a/inference/mlx/trida_mlx/parsing.py b/inference/mlx/trida_mlx/parsing.py new file mode 100644 index 0000000..8230dc6 --- /dev/null +++ b/inference/mlx/trida_mlx/parsing.py @@ -0,0 +1,222 @@ +"""Output parsing for chat responses: ```` reasoning split, Qwen-style tool calls +(JSON ``{...}`` and qwen3-coder XML ````), +and a streaming splitter that routes deltas to reasoning / content / tool-call buffers.""" +from __future__ import annotations + +import json +import re +import uuid + +THINK_OPEN, THINK_CLOSE = "", "" +TC_OPEN, TC_CLOSE = "", "" + + +def _loads_balanced(s: str): + """json.loads the first balanced {...} object in s (tolerates trailing junk like '}}}').""" + depth, start, in_str, esc = 0, None, False, False + for i, ch in enumerate(s): + if in_str: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_str = False + continue + if ch == '"': + in_str = True + elif ch == "{": + if depth == 0: + start = i + depth += 1 + elif ch == "}" and depth > 0: + depth -= 1 + if depth == 0 and start is not None: + try: + return json.loads(s[start : i + 1]) + except Exception: + return None + return None + + +def _param_types(tools) -> dict: + """{tool_name: {param: json-schema type}} from an OpenAI ``tools`` list.""" + out = {} + for t in tools or []: + fn = t.get("function", t) if isinstance(t, dict) else {} + props = ((fn.get("parameters") or {}).get("properties") or {}) if isinstance(fn, dict) else {} + out[fn.get("name")] = {k: (v.get("type") if isinstance(v, dict) else None) for k, v in props.items()} + return out + + +def _coerce(value: str, typ): + """qwen3-coder XML parameters are raw text: keep strings as strings, JSON-decode the rest.""" + if typ == "string" or (isinstance(typ, list) and "string" in typ and len(typ) == 1): + return value + try: + return json.loads(value) + except Exception: + if typ in ("integer", "number"): + try: + return int(value) if typ == "integer" else float(value) + except ValueError: + return value + if typ == "boolean" and value.lower() in ("true", "false"): + return value.lower() == "true" + return value + + +def parse_tool_calls(text: str, tools=None) -> tuple[str, list]: + """Returns (content_without_tool_calls, [{"name", "arguments"(dict)}]).""" + types = _param_types(tools) + calls = [] + for block in re.findall(r"\s*(.*?)\s*(?:|$)", text, re.DOTALL): + obj = _loads_balanced(block) + if isinstance(obj, dict) and "name" in obj: + args = obj.get("arguments", obj.get("parameters", {})) + if isinstance(args, str): + try: + args = json.loads(args) + except Exception: + args = {"_raw": args} + calls.append({"name": obj["name"], "arguments": args if isinstance(args, dict) else {}}) + continue + for fm in re.finditer(r"\s]+)\s*>(.*?)(?:|$)", block, re.DOTALL): + name, a = fm.group(1).strip(), {} + ptypes = types.get(name, {}) + for k, v in re.findall(r"]+?)\s*>\n?(.*?)\n?", fm.group(2), re.DOTALL): + k = k.strip() + typ = ptypes.get(k) + a[k] = _coerce(v if typ == "string" else v.strip(), typ) + calls.append({"name": name, "arguments": a}) + content = re.sub(r".*?(?:|$)", "", text, flags=re.DOTALL).strip() + return content, calls + + +def split_reasoning(text: str, started_in_think: bool) -> tuple[str, str]: + """(reasoning, content). ``started_in_think``: the prompt ended inside an open .""" + if not started_in_think: + if text.lstrip().startswith(THINK_OPEN): + text = text.lstrip()[len(THINK_OPEN):] + started_in_think = True + else: + return "", text + if THINK_CLOSE in text: + r, c = text.split(THINK_CLOSE, 1) + return r.strip(), c.lstrip("\n") + return text.strip(), "" + + +def to_openai_tool_calls(calls: list) -> list: + return [ + { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": {"name": c["name"], "arguments": json.dumps(c["arguments"], ensure_ascii=False)}, + } + for c in calls + ] + + +def _partial_suffix(text: str, tag: str) -> int: + """Length of the longest suffix of text that is a proper prefix of tag.""" + for n in range(min(len(tag) - 1, len(text)), 0, -1): + if tag.startswith(text[-n:]): + return n + return 0 + + +class StreamSplitter: + """Incrementally route generated text into reasoning / content deltas; tool-call blocks + are withheld and returned parsed by ``finish``.""" + + def __init__(self, started_in_think: bool, parse_tools: bool): + self.phase = "maybe_think" if not started_in_think else "think" + self.parse_tools = parse_tools + self.buf = "" + self.tool_text = "" + self.in_tool = False + self.closed_tools = "" + self.content_started = False + + def _content(self, out, text): + if not self.content_started: + text = text.lstrip() + if not text: + return + self.content_started = True + out.append(("content", text)) + + def feed(self, delta: str) -> list[tuple[str, str]]: + """Returns [(kind, text)] with kind in {'reasoning', 'content'}.""" + self.buf += delta + out = [] + while True: + if self.phase == "maybe_think": + s = self.buf.lstrip() + if not s: + return out + if s.startswith(THINK_OPEN): + self.buf = s[len(THINK_OPEN):] + self.phase = "think" + continue + if THINK_OPEN.startswith(s): + return out # could still become + self.phase = "content" + continue + if self.phase == "think": + i = self.buf.find(THINK_CLOSE) + if i >= 0: + if self.buf[:i]: + out.append(("reasoning", self.buf[:i])) + self.buf = self.buf[i + len(THINK_CLOSE):].lstrip("\n") + self.phase = "content" + continue + keep = _partial_suffix(self.buf, THINK_CLOSE) + emit = self.buf[: len(self.buf) - keep] + if emit: + out.append(("reasoning", emit)) + self.buf = self.buf[len(emit):] + return out + # content + if self.in_tool: + self.tool_text += self.buf + self.buf = "" + j = self.tool_text.find(TC_CLOSE) + if j >= 0: + rest = self.tool_text[j + len(TC_CLOSE):] + self.tool_text = self.tool_text[: j + len(TC_CLOSE)] + "\n" + self.in_tool = False + self.buf = rest + self.closed_tools += self.tool_text + self.tool_text = "" + continue + return out + if self.parse_tools: + i = self.buf.find(TC_OPEN) + if i >= 0: + if self.buf[:i]: + self._content(out, self.buf[:i]) + self.in_tool = True + self.tool_text = TC_OPEN + self.buf = self.buf[i + len(TC_OPEN):] + continue + keep = _partial_suffix(self.buf, TC_OPEN) + else: + keep = 0 + emit = self.buf[: len(self.buf) - keep] + if emit: + self._content(out, emit) + self.buf = self.buf[len(emit):] + return out + + def finish(self, tools=None) -> tuple[list[tuple[str, str]], list]: + out = [] + if self.phase == "think" and self.buf: + out.append(("reasoning", self.buf)) + elif self.buf and not self.in_tool: + self._content(out, self.buf) + self.buf = "" + tool_src = self.closed_tools + (self.tool_text if self.in_tool else "") + calls = parse_tool_calls(tool_src, tools)[1] if tool_src else [] + return out, calls diff --git a/inference/mlx/trida_mlx/server.py b/inference/mlx/trida_mlx/server.py new file mode 100644 index 0000000..7c26356 --- /dev/null +++ b/inference/mlx/trida_mlx/server.py @@ -0,0 +1,462 @@ +"""OpenAI-compatible server for on-device Trida (stdlib only: http.server + SSE). + + python -m trida_mlx.server --model ./Trida2.0-4B-mlx-q8 --port 8080 + curl localhost:8080/v1/chat/completions -d '{"messages":[{"role":"user","content":"hi"}]}' + +Endpoints: GET /v1/models, GET /health, GET /stats, POST /v1/chat/completions, +POST /v1/completions. Chat supports ``tools`` (rendered by the model's chat template, +parsed from ```` blocks into OpenAI ``tool_calls``), ``stream``, reasoning split +(``message.reasoning_content``), ``stop``, ``seed`` and per-request +``chat_template_kwargs`` (e.g. ``{"enable_thinking": false}``). Extra body field +``trida_mode`` = ``self-spec`` | ``causal`` overrides the server's decode mode. +One device, one stream: requests are serialized. +""" +from __future__ import annotations + +import argparse +import json +import queue +import sys +import threading +import time +import uuid +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +from .decode import DecodeStats, SamplingParams +from .engine import Engine, IncrementalDetokenizer +from .parsing import StreamSplitter, parse_tool_calls, split_reasoning, to_openai_tool_calls + +ENGINE: Engine = None # set in main() +DEFAULTS: dict = {} + + +def _sampling(body: dict) -> SamplingParams: + t = body.get("temperature", DEFAULTS["temperature"]) + return SamplingParams( + temperature=float(t if t is not None else DEFAULTS["temperature"]), + top_k=int(body.get("top_k") or DEFAULTS["top_k"]), + top_p=float(body.get("top_p") or DEFAULTS["top_p"]), + seed=body.get("seed"), + ) + + +def _text_of(content) -> str: + if isinstance(content, list): + return "".join(p.get("text", "") for p in content if isinstance(p, dict) and "text" in p) + return "" if content is None else str(content) + + +def _normalize_messages(messages: list) -> list: + """OpenAI -> the Trida (Qwen3.5) chat template's contract: + * content parts -> text (no vision here), None -> "" + * ``developer`` -> ``system``; all leading system messages merged into one; a system message + later in the conversation (harness nudges, compaction notes) becomes a user turn, because + the template only accepts a system message first + * assistant ``tool_calls[].function.arguments`` as dicts (the template iterates them) + """ + out, sys_parts, seen_non_system = [], [], False + for m in messages: + m = dict(m) + role = m.get("role") + if role == "developer": + role = m["role"] = "system" + m["content"] = _text_of(m.get("content")) + if role == "system": + if not seen_non_system: + sys_parts.append(m["content"]) + continue + m = {"role": "user", "content": f"[System note]\n{m['content']}"} + seen_non_system = True + if role == "assistant": + if isinstance(m.get("reasoning"), str) and not isinstance(m.get("reasoning_content"), str): + m["reasoning_content"] = m["reasoning"] + if m.get("tool_calls"): + tcs = [] + for tc in m["tool_calls"]: + fn = dict(tc.get("function", tc)) + args = fn.get("arguments") + if isinstance(args, str): + try: + args = json.loads(args) if args.strip() else {} + except Exception: + args = {"_raw": args} + fn["arguments"] = args if isinstance(args, dict) else {"value": args} + tcs.append({"type": "function", "id": tc.get("id"), "function": fn}) + m["tool_calls"] = tcs + out.append(m) + if sys_parts: + out.insert(0, {"role": "system", "content": "\n\n".join(p for p in sys_parts if p)}) + return out + + +def _thinking(body: dict) -> bool: + kw = body.get("chat_template_kwargs") or {} + if "enable_thinking" in kw: + return bool(kw["enable_thinking"]) + if "enable_thinking" in body: + return bool(body["enable_thinking"]) + eff = body.get("reasoning_effort") + if eff is None and isinstance(body.get("reasoning"), dict): + eff = body["reasoning"].get("effort") + if isinstance(eff, str) and eff.lower() in ("none", "off", "disabled", "minimal"): + return False + return DEFAULTS["enable_thinking"] + + +class ContextOverflow(Exception): + pass + + +class _StopMatcher: + def __init__(self, stops): + self.stops = [s for s in ([stops] if isinstance(stops, str) else (stops or [])) if s] + self.text = "" + + def check(self, full_text: str): + """Returns index to cut the full text at, or None.""" + for s in self.stops: + i = full_text.find(s) + if i >= 0: + return i + return None + + +class _Job: + def __init__(self, prompt, body, max_tokens, want_progress): + self.prompt, self.body, self.max_tokens, self.want_progress = prompt, body, max_tokens, want_progress + self.out: queue.Queue = queue.Queue() + self.cancelled = False + + def run(self): + """Runs on the worker thread (the only thread that touches MLX).""" + body = self.body + stop = _StopMatcher(body.get("stop")) + stats = DecodeStats() + text = "" + detok = IncrementalDetokenizer(ENGINE.tokenizer) + progress = (lambda done, total: self.out.put(("progress", done, total))) if self.want_progress else None + gen = ENGINE.generate_ids(self.prompt, max_tokens=self.max_tokens, sampling=_sampling(body), + mode=body.get("trida_mode"), stats=stats, progress=progress) + try: + for ids in gen: + if self.cancelled: + stats.finish_reason = "cancelled" + break + d = detok.add(ids) + if not d: + continue + cut = stop.check(text + d) if stop.stops else None + if cut is not None: + d = (text + d)[len(text):cut] if cut >= len(text) else "" + text += d + if d: + self.out.put(("delta", d)) + stats.finish_reason = "stop" + break + text += d + self.out.put(("delta", d)) + finally: + gen.close() + d = stats.as_dict() + sys.stderr.write(f"[trida-mlx] prompt {d['prompt_tokens']} (cached {d['reused_tokens']}, prefill " + f"{d['prefill_s']:.1f}s) -> {d['new_tokens']} tok @ {d['decode_tok_s']} tok/s, " + f"{d['tokens_per_forward']} tok/fwd [{d['finish_reason']}]\n") + self.out.put(("done", stats)) + + +class _Worker(threading.Thread): + """MLX streams are per-thread, so the model is loaded and every forward runs on this one + thread; HTTP handler threads only render/tokenize and relay deltas.""" + + def __init__(self, engine_kwargs: dict): + super().__init__(daemon=True, name="trida-mlx-worker") + self.engine_kwargs = engine_kwargs + self.jobs: queue.Queue = queue.Queue() + self.ready = threading.Event() + self.error = None + + def submit(self, job: "_Job"): + self.jobs.put(job) + + def run(self): + global ENGINE + try: + ENGINE = Engine(**self.engine_kwargs) + ENGINE.warmup() + except BaseException as e: # noqa: BLE001 + self.error = e + self.ready.set() + return + self.ready.set() + while True: + job = self.jobs.get() + if job.cancelled: + job.out.put(("done", DecodeStats(finish_reason="cancelled"))) + continue + try: + job.run() + except BaseException as e: # noqa: BLE001 + import traceback + traceback.print_exc() + job.out.put(("error", e)) + + +WORKER: "_Worker" = None + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, fmt, *args): # quieter + sys.stderr.write("[trida-mlx] " + (fmt % args) + "\n") + + # -- helpers ---------------------------------------------------------------- + def _json(self, code: int, obj: dict): + data = json.dumps(obj, ensure_ascii=False).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(data) + + def _sse_start(self): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "close") + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.close_connection = True + + def _sse(self, obj): + payload = obj if isinstance(obj, str) else json.dumps(obj, ensure_ascii=False) + self.wfile.write(f"data: {payload}\n\n".encode()) + self.wfile.flush() + + def do_OPTIONS(self): + self.send_response(204) + self.send_header("Access-Control-Allow-Origin", "*") + self.send_header("Access-Control-Allow-Headers", "*") + self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") + self.send_header("Content-Length", "0") + self.end_headers() + + def _model_card(self): + ctx = DEFAULTS["context_length"] + return {"id": DEFAULTS["served_name"], "object": "model", "owned_by": "trillionlabs", "created": 0, + "context_length": ctx, "max_model_len": ctx, "root": str(ENGINE.path)} + + def do_GET(self): + path = self.path.split("?")[0].rstrip("/") + if path in ("/v1/models", "/models"): + return self._json(200, {"object": "list", "data": [self._model_card()]}) + if path.startswith(("/v1/models/", "/models/")): + return self._json(200, self._model_card()) + if path in ("/health", "/v1/health", ""): + return self._json(200, {"status": "ok", "mode": ENGINE.mode}) + if path == "/stats": + return self._json(200, {"last": ENGINE.last_stats, "cache_slots": ENGINE.cached_tokens()}) + self._json(404, {"error": {"message": f"no route {self.path}"}}) + + def do_POST(self): + try: + n = int(self.headers.get("Content-Length") or 0) + body = json.loads(self.rfile.read(n) or b"{}") + except Exception as e: + return self._json(400, {"error": {"message": f"bad json: {e}"}}) + route = self.path.split("?")[0].rstrip("/") + if DEFAULTS.get("log_requests"): + import pathlib + d = pathlib.Path(DEFAULTS["log_requests"]); d.mkdir(parents=True, exist_ok=True) + (d / f"{time.strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:6]}.json").write_text( + json.dumps({"path": route, "body": body}, ensure_ascii=False, indent=1)) + try: + if route in ("/v1/chat/completions", "/chat/completions"): + return self._chat(body) + if route in ("/v1/completions", "/completions"): + return self._completion(body) + return self._json(404, {"error": {"message": f"no route {self.path}"}}) + except ContextOverflow as e: + return self._json(400, {"error": {"message": str(e), "type": "invalid_request_error", + "param": "messages", "code": "context_length_exceeded"}}) + except (BrokenPipeError, ConnectionResetError): + pass + except Exception as e: # noqa: BLE001 + import traceback + traceback.print_exc() + try: + self._json(500, {"error": {"message": f"{type(e).__name__}: {e}"}}) + except Exception: + pass + + # -- generation core -------------------------------------------------------- + @staticmethod + def _budget(prompt: list, body: dict) -> int: + ctx = DEFAULTS["context_length"] + if len(prompt) >= ctx: + raise ContextOverflow( + f"This model's maximum context length is {ctx} tokens. However, your messages resulted in " + f"{len(prompt)} tokens. Please reduce the length of the messages.") + want = int(body.get("max_completion_tokens") or body.get("max_tokens") or DEFAULTS["max_tokens"]) + return max(1, min(want, ctx - len(prompt))) + + def _run(self, prompt: list, body: dict, on_delta, keepalive=None): + """Generate on the MLX worker thread; call on_delta(text) here per decoded delta. + Returns (full_text, stats). A client disconnect cancels the generation.""" + max_tokens = self._budget(prompt, body) + job = _Job(prompt, body, max_tokens, want_progress=keepalive is not None) + WORKER.submit(job) + text = "" + try: + while True: + kind, *val = job.out.get() + if kind == "delta": + text += val[0] + on_delta(val[0]) + elif kind == "progress": + keepalive(*val) + elif kind == "error": + raise val[0] + elif kind == "done": + return text, val[0] + except BaseException: + job.cancelled = True + raise + + def _usage(self, stats: DecodeStats): + return {"prompt_tokens": stats.prompt_tokens, "completion_tokens": stats.new_tokens, + "total_tokens": stats.prompt_tokens + stats.new_tokens, + "prompt_tokens_details": {"cached_tokens": stats.reused_tokens}, + "trida": stats.as_dict()} + + def _chat(self, body: dict): + messages = _normalize_messages(body.get("messages") or []) + tools = body.get("tools") if body.get("tool_choice") != "none" else None + kw = {k: v for k, v in (body.get("chat_template_kwargs") or {}).items() if k != "enable_thinking"} + prompt_text = ENGINE.render(messages, tools=tools, enable_thinking=_thinking(body), **kw) + prompt = ENGINE.encode(prompt_text) + self._budget(prompt, body) # 400 before any streaming starts + started_in_think = prompt_text.rstrip().endswith("") + rid, created = f"chatcmpl-{uuid.uuid4().hex[:24]}", int(time.time()) + model = body.get("model") or DEFAULTS["served_name"] + + if not body.get("stream"): + text, stats = self._run(prompt, body, lambda d: None) + reasoning, content = split_reasoning(text, started_in_think) + calls = [] + if tools: + content, calls = parse_tool_calls(content, tools) + msg = {"role": "assistant", "content": content or (None if calls else "")} + if reasoning: + msg["reasoning_content"] = reasoning + finish = stats.finish_reason + if calls: + msg["tool_calls"] = to_openai_tool_calls(calls) + finish = "tool_calls" + return self._json(200, {"id": rid, "object": "chat.completion", "created": created, "model": model, + "choices": [{"index": 0, "message": msg, "finish_reason": finish}], + "usage": self._usage(stats)}) + + self._sse_start() + splitter = StreamSplitter(started_in_think, parse_tools=bool(tools)) + + def chunk(delta: dict, finish=None, **extra): + obj = {"id": rid, "object": "chat.completion.chunk", "created": created, "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]} + obj.update(extra) + self._sse(obj) + + def emit(parts): + for kind, t in parts: + chunk({"reasoning_content": t} if kind == "reasoning" else {"content": t}) + + def keepalive(done, total): + self.wfile.write(f": prefill {done}/{total}\n\n".encode()) + self.wfile.flush() + + chunk({"role": "assistant", "content": ""}) + _, stats = self._run(prompt, body, lambda d: emit(splitter.feed(d)), keepalive=keepalive) + rest, calls = splitter.finish(tools) + emit(rest) + finish = stats.finish_reason + if calls: + tcs = to_openai_tool_calls(calls) + chunk({"tool_calls": [dict(tc, index=i) for i, tc in enumerate(tcs)]}) + finish = "tool_calls" + chunk({}, finish) + if (body.get("stream_options") or {}).get("include_usage", True): + self._sse({"id": rid, "object": "chat.completion.chunk", "created": created, "model": model, + "choices": [], "usage": self._usage(stats)}) + self._sse("[DONE]") + + def _completion(self, body: dict): + prompt_text = body.get("prompt") or "" + if isinstance(prompt_text, list): + prompt_text = prompt_text[0] + rid, created = f"cmpl-{uuid.uuid4().hex[:24]}", int(time.time()) + model = body.get("model") or DEFAULTS["served_name"] + prompt = ENGINE.encode(prompt_text) + self._budget(prompt, body) + if not body.get("stream"): + text, stats = self._run(prompt, body, lambda d: None) + return self._json(200, {"id": rid, "object": "text_completion", "created": created, "model": model, + "choices": [{"index": 0, "text": text, "finish_reason": stats.finish_reason}], + "usage": self._usage(stats)}) + self._sse_start() + _, stats = self._run(prompt, body, lambda d: self._sse( + {"id": rid, "object": "text_completion", "created": created, "model": model, + "choices": [{"index": 0, "text": d, "finish_reason": None}]})) + self._sse({"id": rid, "object": "text_completion", "created": created, "model": model, + "choices": [{"index": 0, "text": "", "finish_reason": stats.finish_reason}], + "usage": self._usage(stats)}) + self._sse("[DONE]") + + +def main(argv=None): + global WORKER + ap = argparse.ArgumentParser(description="OpenAI-compatible on-device Trida server (MLX)") + ap.add_argument("--model", default="trillionlabs/Trida2.0-4B", + help="HF id or local dir (raw checkpoint or convert.py output)") + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port", type=int, default=8080) + ap.add_argument("--mode", default="self-spec", choices=["self-spec", "causal"]) + ap.add_argument("--gen-block", type=int, default=4, help="N: tokens per self-spec step (canvas 2N-1)") + ap.add_argument("--prefill-chunk", type=int, default=512) + ap.add_argument("--no-prompt-cache", action="store_true") + ap.add_argument("--temperature", type=float, default=1.0) + ap.add_argument("--top-k", type=int, default=50) + ap.add_argument("--top-p", type=float, default=0.95) + ap.add_argument("--max-tokens", type=int, default=8192, help="default output budget (capped by context)") + ap.add_argument("--context-length", type=int, default=65536, + help="advertised + enforced context window (Hermes needs >= 64000)") + ap.add_argument("--served-model-name", default="trida2.0-4b") + ap.add_argument("--log-requests", default=None, metavar="DIR", help="dump every POST body as JSON (debugging)") + ap.add_argument("--cache-slots", type=int, default=3, + help="conversations whose prompt cache is kept (main loop + side requests)") + ap.add_argument("--no-think", action="store_true", help="default enable_thinking=false") + ap.add_argument("--no-fused-gdn", action="store_true") + a = ap.parse_args(argv) + DEFAULTS.update(temperature=a.temperature, top_k=a.top_k, top_p=a.top_p, max_tokens=a.max_tokens, + enable_thinking=not a.no_think, context_length=a.context_length, + served_name=a.served_model_name, log_requests=a.log_requests) + print(f"[trida-mlx] loading {a.model} ...", flush=True) + WORKER = _Worker(dict(model=a.model, mode=a.mode, gen_block=a.gen_block, prefill_chunk=a.prefill_chunk, + prompt_cache=not a.no_prompt_cache, fused_gdn=not a.no_fused_gdn, + cache_slots=a.cache_slots)) + WORKER.start() + WORKER.ready.wait() + if WORKER.error is not None: + raise WORKER.error + print(f"[trida-mlx] ready in {ENGINE.load_s:.1f}s mode={a.mode} N={a.gen_block} mask_id={ENGINE.mask_id} " + f"eos={sorted(ENGINE.eos_ids)} ctx={a.context_length} model={a.served_model_name!r} -> " + f"http://{a.host}:{a.port}/v1", flush=True) + srv = ThreadingHTTPServer((a.host, a.port), Handler) + try: + srv.serve_forever() + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() From 9d20230968089cca60cd55b3dfaf1ff97062188f Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Wed, 30 Sep 2026 01:29:25 +0000 Subject: [PATCH 3/9] docs(inference/mlx): document the on-device backend and its third-party code - inference/mlx/README.md: quickstart (uv), design mapping to the SGLang reference, prompt cache, Hermes setup, measured M3 Pro numbers, limits - root README layout + CHANGELOG entry - NOTICE / COMPLIANCE: the mlx-lm-derived Metal kernel (MIT) and the bundled chat-template test fixture (Apache-2.0); MLX + mlx-lm as runtime dependencies Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- CHANGELOG.md | 4 ++ COMPLIANCE.md | 6 +- NOTICE | 20 +++++- README.md | 1 + inference/mlx/README.md | 148 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 inference/mlx/README.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a6cf8d..6bbef3a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- `inference/mlx/`: on-device (Apple Silicon, MLX) serving of Trida2.0-4B with lossless + self-speculative decoding (`bd_bidir_shift` b7/g4 semantics), an OpenAI-compatible server with + tool calls and a prompt cache, a minimal agent loop, MLX quantization, and offline tests. + ### Changed ### Fixed diff --git a/COMPLIANCE.md b/COMPLIANCE.md index bc692af..6c256eb 100644 --- a/COMPLIANCE.md +++ b/COMPLIANCE.md @@ -11,8 +11,10 @@ model-card license. | Component | Path | Upstream | License | |---|---|---|---| | IFEval scorer | `inference/ifeval_lib/` | [google-research · instruction_following_eval](https://github.com/google-research/google-research/tree/master/instruction_following_eval) | **Apache-2.0** | +| MLX gated-delta Metal kernel (adapted) | `inference/mlx/trida_mlx/kernels.py` | [ml-explore/mlx-lm](https://github.com/ml-explore/mlx-lm) `mlx_lm/models/gated_delta.py` | **MIT** | +| Qwen3.5 / Trida chat template (test fixture) | `inference/mlx/tests/fixtures/trida_chat_template.jinja` | shipped with `trillionlabs/Trida2.0-4B` (from Qwen3.5) | **Apache-2.0** | -This is the **only** third-party source code redistributed in the repository (see +These are the only third-party source code redistributed in the repository (see [`NOTICE`](NOTICE)). Everything else below is installed or fetched separately. ## 1b. Noncommercial upstream code (fetched via recipe, NOT bundled) @@ -43,6 +45,8 @@ noncommercial code. | HF Accelerate | training | [huggingface/accelerate](https://github.com/huggingface/accelerate) | **Apache-2.0** | | SGLang | diffusion serving backend | [sgl-project/sglang](https://github.com/sgl-project/sglang) | **Apache-2.0** | | vLLM | optional AR serving backend | [vllm-project/vllm](https://github.com/vllm-project/vllm) | **Apache-2.0** | +| MLX | on-device backend (`inference/mlx/`) | [ml-explore/mlx](https://github.com/ml-explore/mlx) | **MIT** | +| mlx-lm | Qwen3.5 model + kernels for `inference/mlx/` | [ml-explore/mlx-lm](https://github.com/ml-explore/mlx-lm) | **MIT** | ## 3. Benchmark harnesses (invoked from `benchmark/`, not bundled) diff --git a/NOTICE b/NOTICE index 4d6f3f8..3807748 100644 --- a/NOTICE +++ b/NOTICE @@ -31,7 +31,22 @@ Two-stream Gated-DeltaNet / ShortConv Triton kernels Note: inference/sglang/ vendors nothing — it builds the backend from upstream via a patch; see inference/sglang/NOTICE for that recipe's terms. -The two components above are the third-party source code redistributed in this +MLX gated-delta Metal kernel (mlx-lm) + Copyright © 2023-2026 Apple Inc. + License: MIT + Source: https://github.com/ml-explore/mlx-lm (mlx_lm/models/gated_delta.py) + Derived: inference/mlx/trida_mlx/kernels.py — the self-spec canvas kernel adapts + mlx-lm's gated_delta Metal kernel source (same thread layout and + recurrence; adds the block-end readout for MASK rows). See its header. + +Qwen3.5 chat template (test fixture) + Copyright Alibaba Cloud (Qwen team). + License: Apache License 2.0 + Bundled: inference/mlx/tests/fixtures/trida_chat_template.jinja — the chat template + shipped with trillionlabs/Trida2.0-4B (inherited from Qwen3.5), unmodified, + used only by the offline tests. + +The components above are the third-party source code redistributed in this repository. They do not share a single license: see COMPLIANCE.md section 1. ================================================================================ @@ -58,7 +73,8 @@ Referenced, NOT bundled (installed separately) ================================================================================ Runtime backends and libraries — PyTorch, Hugging Face Transformers / Datasets / -Accelerate, SGLang, and optionally vLLM — and the benchmark harnesses invoked from +Accelerate, SGLang, optionally vLLM, and for inference/mlx Apple MLX and mlx-lm +(MIT) — and the benchmark harnesses invoked from benchmark/ (EleutherAI lm-evaluation-harness, BFCL, tau2-bench, FunctionChat, KoAgentBench, SWE-bench, Terminal-Bench) are dependencies, not redistributed here. See COMPLIANCE.md for the full inventory and licenses. diff --git a/README.md b/README.md index 7455d5e..d96e318 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,7 @@ inference/ nano-inference serving stack (thin, nanoGPT-style) vllm/ second backend: out-of-tree vLLM plugin (block-diffusion + self-spec/AR-Trust) vllm_native_diffusion/ the plugin package (VLLM_PLUGINS=trida_diffusion) tools/, docs/ Slurm eval jobs, worklog + design notes + mlx/ on-device backend (Apple Silicon): MLX self-spec decoder, OpenAI server, agent benchmark/ heavier agentic eval harnesses (bfcl_v4, tau2, functionchat, ko_agentbench, swe_bench, terminal_bench) + serving/ data/ dataset download scripts + catalog diff --git a/inference/mlx/README.md b/inference/mlx/README.md new file mode 100644 index 0000000..018598f --- /dev/null +++ b/inference/mlx/README.md @@ -0,0 +1,148 @@ +# inference/mlx — on-device self-speculative decoding (Apple Silicon) + +Runs **[`trillionlabs/Trida2.0-4B`](https://huggingface.co/trillionlabs/Trida2.0-4B)** on a Mac with +**MLX**, in the same three pieces as the GPU stacks next door — but in-process, one stream, no CUDA: + +| | file | what | +|---|---|---| +| decoder | `trida_mlx/model.py`, `trida_mlx/decode.py` | `causal` (AR baseline) and **`self-spec`** (diffusion draft + exact AR verify, lossless) | +| server | `trida_mlx/server.py` | OpenAI-compatible `/v1/chat/completions` (+ `/v1/completions`), SSE streaming, **tool calls**, reasoning split, prompt cache | +| agent | `trida_mlx/agent.py` | a tiny tool-using agent loop against the server (files, search, calculator; opt-in shell/write) | +| tools | `convert.py`, `verify.py`, `bench.py` | quantize to MLX, first-contact correctness check, causal-vs-self-spec benchmark | + +Everything is built on `mlx_lm.models.qwen3_5` (weights, quantized layers, RoPE and the Metal +gated-delta kernel come from mlx-lm unchanged); this package only adds the self-spec forward, +the cache commit logic and the serving layer. It follows the approach of the `ocr-monorepo` +MLX self-spec OCR decoder, extended to the **hybrid** Qwen3.5 backbone. + +## Quickstart + +Run these one by one (zsh does not treat `#` as a comment in interactive shells, so don't paste +trailing comments). If another venv is active, `deactivate` first. + +```bash +cd inference/mlx +uv sync +uv run hf auth login +uv run trida-mlx-convert --model trillionlabs/Trida2.0-4B --out ~/models/Trida2.0-4B-mlx-q8 --bits 8 +uv run trida-mlx-verify --model ~/models/Trida2.0-4B-mlx-q8 +uv run trida-mlx-bench --model ~/models/Trida2.0-4B-mlx-q8 --max-tokens 512 --out bench_q8.json +uv run trida-mlx-server --model ~/models/Trida2.0-4B-mlx-q8 --port 8080 +uv run trida-mlx-agent --workdir ~/some/project "summarize what this repo does" +``` + +The console scripts are aliases for `python -m trida_mlx.`. +Without uv: `pip install -e .` (or `pip install -r requirements.txt` and run from this directory). + +`uv sync` makes a Python 3.12 `.venv` from `uv.lock`. `convert` resolves the private checkpoint +(needs `hf auth login`), `verify` checks canvas rows == AR logits and greedy self-spec == greedy AR, +`bench` compares causal vs self-spec (add `--temperature 1` for sampled), then serve and point the +agent (second terminal) at it. q8 is ≈4.5 GB, bf16 ≈8 GB; on 16–18 GB Macs use q8 or q6. + +Any OpenAI client works against the server: + +```python +from openai import OpenAI +c = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="local") +r = c.chat.completions.create(model="trida", messages=[{"role": "user", "content": "hi"}], + extra_body={"chat_template_kwargs": {"enable_thinking": False}}) +print(r.choices[0].message.content, r.usage) # usage.trida has tok/s, tok/forward, accept histogram +``` + +Server flags: `--mode self-spec|causal`, `--gen-block N` (canvas `2N-1`, default 4 → the shipped +`b7/g4` config), `--temperature/--top-k/--top-p` defaults (1.0 / 50 / 0.95, the canonical Trida +self-spec setting), `--no-think`, `--max-tokens`, `--no-prompt-cache`. Per request: `trida_mode`, +`chat_template_kwargs`, `tools`, `stop`, `seed`, `stream`. + +## How self-spec maps onto MLX + +Reference semantics are the SGLang `HybridDiffusionSelfSpec` (`bd_bidir_shift`, block 7, gen 4, +`causal_prefill`, `draft_mode: strict_truncated`) — the same algorithm the vLLM plugin ports. + +``` +cold start : [t0, M, M, M, M, M, M ] clean = row 0, specs = rows 1..3 +verify : [p, s0, s1, s2, M, M, M ] s_i checked against row i (exact AR logits) + all accepted -> 4 tokens: s0 s1 s2 + clean(row 3), new specs = rows 4..6 + reject at i -> s_` of the previous prompt (where the next +turn's re-rendered history still agrees with it). A follow-up turn then only prefills the new messages; `usage.prompt_tokens_details.cached_tokens` +shows the reuse. + +## Tests + +```bash +uv run pytest # all offline tests (~30 s) +uv run python tests/test_tiny.py # decoder invariants on a tiny random Qwen3.5-hybrid (CPU or GPU) +uv run python tests/test_e2e.py # tiny raw-HF checkpoint: loader, convert q8, prompt cache, server, parsing +``` + +Pinned: chunked prefill = one-shot; canvas clean rows = AR logits; MASK rows read the block-end GDN +state (independent per-token re-derivation); `commit(adv)` = AR state for every `adv`; greedy +self-spec = greedy AR (also with oracle drafts that exercise the accept path, ~3 tok/forward); +top-k/top-p matches the reference sampler; sampled self-spec reproduces the AR sampling +distribution; snapshot/restore = fresh prefill; server non-stream / SSE / completions round trip. +On a Mac the same tests run on the GPU and therefore also cover the Metal kernel path. + +## Notes and limits + +- **bf16/quantized numerics.** The canvas runs 7-row matmuls and the AR step 1-row matmuls, which + are different kernels; at a near-tie the greedy argmax can flip. `verify.py` prints the AR + top-1/top-2 margin at the first divergence so you can tell numerics (tiny margin) from a bug. +- **Quantization.** Self-spec is lossless relative to *the model it runs*; quantization changes that + model. In the OCR study q8/q6 kept outputs identical to bf16 and q4 did not; draft acceptance can + also drop at low bit widths. `convert.py --keep-embed` keeps the tied LM head in bf16. +- **Single stream.** Requests are serialized (one device); no batching. +- **Step cost.** `trida-mlx-profile` splits the AR step and the canvas step into LM head / MLPs / + attention / gated-delta and prints the break-even tokens-per-forward. Done: a fused canvas GDN + Metal kernel (`trida_mlx/kernels.py`, one launch per layer; `--no-fused-gdn` to A/B): verify step + 49.8 → 44.5 ms on M3 Pro q8, i.e. 1.13× an AR step (39.3 ms). LM-head row skipping on cold starts + is available (`bench --skip-rows`) but measured slower, so it is off. +- The GGUF quants (`mradermacher/Trida2.0-4B-*GGUF`) run in llama.cpp as plain AR only; self-spec + there would need the canvas mask and the block-end GDN readout in llama.cpp's qwen35 graph. From d344fdcd63993ae563b7ea652c65171a5138b541 Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Wed, 30 Sep 2026 09:21:58 +0000 Subject: [PATCH 4/9] chore(inference/mlx): drop the tests directory Remove inference/mlx/tests (tiny-model decoder/server tests and the chat template fixture) and everything that referenced it: the README Tests section, the pytest dev group and config (uv.lock re-locked), and the fixture's NOTICE / COMPLIANCE entries. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- CHANGELOG.md | 2 +- COMPLIANCE.md | 1 - NOTICE | 7 - inference/mlx/README.md | 15 - inference/mlx/pyproject.toml | 9 +- .../tests/fixtures/trida_chat_template.jinja | 154 --------- inference/mlx/tests/test_e2e.py | 229 -------------- inference/mlx/tests/test_tiny.py | 297 ------------------ inference/mlx/tests/tiny_ckpt.py | 96 ------ inference/mlx/uv.lock | 47 +-- 10 files changed, 4 insertions(+), 853 deletions(-) delete mode 100644 inference/mlx/tests/fixtures/trida_chat_template.jinja delete mode 100644 inference/mlx/tests/test_e2e.py delete mode 100644 inference/mlx/tests/test_tiny.py delete mode 100644 inference/mlx/tests/tiny_ckpt.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1667828..016345b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `inference/mlx/`: on-device (Apple Silicon, MLX) serving of Trida2.0-4B with lossless self-speculative decoding (`bd_bidir_shift` b7/g4 semantics), an OpenAI-compatible server with - tool calls and a prompt cache, a minimal agent loop, MLX quantization, and offline tests. + tool calls and a prompt cache, a minimal agent loop, and MLX quantization. ## [0.1.0] — 2026-09-30 diff --git a/COMPLIANCE.md b/COMPLIANCE.md index 6c256eb..ae12c2c 100644 --- a/COMPLIANCE.md +++ b/COMPLIANCE.md @@ -12,7 +12,6 @@ model-card license. |---|---|---|---| | IFEval scorer | `inference/ifeval_lib/` | [google-research · instruction_following_eval](https://github.com/google-research/google-research/tree/master/instruction_following_eval) | **Apache-2.0** | | MLX gated-delta Metal kernel (adapted) | `inference/mlx/trida_mlx/kernels.py` | [ml-explore/mlx-lm](https://github.com/ml-explore/mlx-lm) `mlx_lm/models/gated_delta.py` | **MIT** | -| Qwen3.5 / Trida chat template (test fixture) | `inference/mlx/tests/fixtures/trida_chat_template.jinja` | shipped with `trillionlabs/Trida2.0-4B` (from Qwen3.5) | **Apache-2.0** | These are the only third-party source code redistributed in the repository (see [`NOTICE`](NOTICE)). Everything else below is installed or fetched separately. diff --git a/NOTICE b/NOTICE index 3807748..fea11f5 100644 --- a/NOTICE +++ b/NOTICE @@ -39,13 +39,6 @@ MLX gated-delta Metal kernel (mlx-lm) mlx-lm's gated_delta Metal kernel source (same thread layout and recurrence; adds the block-end readout for MASK rows). See its header. -Qwen3.5 chat template (test fixture) - Copyright Alibaba Cloud (Qwen team). - License: Apache License 2.0 - Bundled: inference/mlx/tests/fixtures/trida_chat_template.jinja — the chat template - shipped with trillionlabs/Trida2.0-4B (inherited from Qwen3.5), unmodified, - used only by the offline tests. - The components above are the third-party source code redistributed in this repository. They do not share a single license: see COMPLIANCE.md section 1. diff --git a/inference/mlx/README.md b/inference/mlx/README.md index 018598f..9d9000f 100644 --- a/inference/mlx/README.md +++ b/inference/mlx/README.md @@ -115,21 +115,6 @@ system block, and one just before the last `<|im_start|>` of the previous prompt turn's re-rendered history still agrees with it). A follow-up turn then only prefills the new messages; `usage.prompt_tokens_details.cached_tokens` shows the reuse. -## Tests - -```bash -uv run pytest # all offline tests (~30 s) -uv run python tests/test_tiny.py # decoder invariants on a tiny random Qwen3.5-hybrid (CPU or GPU) -uv run python tests/test_e2e.py # tiny raw-HF checkpoint: loader, convert q8, prompt cache, server, parsing -``` - -Pinned: chunked prefill = one-shot; canvas clean rows = AR logits; MASK rows read the block-end GDN -state (independent per-token re-derivation); `commit(adv)` = AR state for every `adv`; greedy -self-spec = greedy AR (also with oracle drafts that exercise the accept path, ~3 tok/forward); -top-k/top-p matches the reference sampler; sampled self-spec reproduces the AR sampling -distribution; snapshot/restore = fresh prefill; server non-stream / SSE / completions round trip. -On a Mac the same tests run on the GPU and therefore also cover the Metal kernel path. - ## Notes and limits - **bf16/quantized numerics.** The canvas runs 7-row matmuls and the AR step 1-row matmuls, which diff --git a/inference/mlx/pyproject.toml b/inference/mlx/pyproject.toml index 2b6714e..b8f761b 100644 --- a/inference/mlx/pyproject.toml +++ b/inference/mlx/pyproject.toml @@ -11,7 +11,7 @@ requires-python = ">=3.10,<3.14" license = { text = "Apache-2.0" } authors = [{ name = "Trillion Labs" }] dependencies = [ - # Metal build on macOS; the CPU build elsewhere (Linux CI / offline tests only). + # Metal build on macOS; the CPU build elsewhere (Linux, no GPU). "mlx>=0.29; sys_platform == 'darwin'", "mlx[cpu]>=0.29; sys_platform == 'linux'", "mlx-lm>=0.31", # Qwen3.5 gated-delta model + Metal kernel we build on @@ -27,16 +27,9 @@ trida-mlx-verify = "trida_mlx.verify:main" trida-mlx-bench = "trida_mlx.bench:main" trida-mlx-profile = "trida_mlx.profile:main" -[dependency-groups] -dev = ["pytest>=8"] - [tool.setuptools] packages = ["trida_mlx"] [tool.uv] # resolve only for the platforms this package runs on environments = ["sys_platform == 'darwin' and platform_machine == 'arm64'", "sys_platform == 'linux'"] - -[tool.pytest.ini_options] -testpaths = ["tests"] -pythonpath = [".", "tests"] diff --git a/inference/mlx/tests/fixtures/trida_chat_template.jinja b/inference/mlx/tests/fixtures/trida_chat_template.jinja deleted file mode 100644 index ade1daa..0000000 --- a/inference/mlx/tests/fixtures/trida_chat_template.jinja +++ /dev/null @@ -1,154 +0,0 @@ -{%- set image_count = namespace(value=0) %} -{%- set video_count = namespace(value=0) %} -{%- macro render_content(content, do_vision_count, is_system_content=false) %} - {%- if content is string %} - {{- content }} - {%- elif content is iterable and content is not mapping %} - {%- for item in content %} - {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} - {%- if is_system_content %} - {{- raise_exception('System message cannot contain images.') }} - {%- endif %} - {%- if do_vision_count %} - {%- set image_count.value = image_count.value + 1 %} - {%- endif %} - {%- if add_vision_id %} - {{- 'Picture ' ~ image_count.value ~ ': ' }} - {%- endif %} - {{- '<|vision_start|><|image_pad|><|vision_end|>' }} - {%- elif 'video' in item or item.type == 'video' %} - {%- if is_system_content %} - {{- raise_exception('System message cannot contain videos.') }} - {%- endif %} - {%- if do_vision_count %} - {%- set video_count.value = video_count.value + 1 %} - {%- endif %} - {%- if add_vision_id %} - {{- 'Video ' ~ video_count.value ~ ': ' }} - {%- endif %} - {{- '<|vision_start|><|video_pad|><|vision_end|>' }} - {%- elif 'text' in item %} - {{- item.text }} - {%- else %} - {{- raise_exception('Unexpected item type in content.') }} - {%- endif %} - {%- endfor %} - {%- elif content is none or content is undefined %} - {{- '' }} - {%- else %} - {{- raise_exception('Unexpected content type.') }} - {%- endif %} -{%- endmacro %} -{%- if not messages %} - {{- raise_exception('No messages provided.') }} -{%- endif %} -{%- if tools and tools is iterable and tools is not mapping %} - {{- '<|im_start|>system\n' }} - {{- "# Tools\n\nYou have access to the following functions:\n\n" }} - {%- for tool in tools %} - {{- "\n" }} - {{- tool | tojson }} - {%- endfor %} - {{- "\n" }} - {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} - {%- if messages[0].role == 'system' %} - {%- set content = render_content(messages[0].content, false, true)|trim %} - {%- if content %} - {{- '\n\n' + content }} - {%- endif %} - {%- endif %} - {{- '<|im_end|>\n' }} -{%- else %} - {%- if messages[0].role == 'system' %} - {%- set content = render_content(messages[0].content, false, true)|trim %} - {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} - {%- endif %} -{%- endif %} -{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} -{%- for message in messages[::-1] %} - {%- set index = (messages|length - 1) - loop.index0 %} - {%- if ns.multi_step_tool and message.role == "user" %} - {%- set content = render_content(message.content, false)|trim %} - {%- if not(content.startswith('') and content.endswith('')) %} - {%- set ns.multi_step_tool = false %} - {%- set ns.last_query_index = index %} - {%- endif %} - {%- endif %} -{%- endfor %} -{%- if ns.multi_step_tool %} - {{- raise_exception('No user query found in messages.') }} -{%- endif %} -{%- for message in messages %} - {%- set content = render_content(message.content, true)|trim %} - {%- if message.role == "system" %} - {%- if not loop.first %} - {{- raise_exception('System message must be at the beginning.') }} - {%- endif %} - {%- elif message.role == "user" %} - {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} - {%- elif message.role == "assistant" %} - {%- set reasoning_content = '' %} - {%- if message.reasoning_content is string %} - {%- set reasoning_content = message.reasoning_content %} - {%- else %} - {%- if '' in content %} - {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} - {%- set content = content.split('')[-1].lstrip('\n') %} - {%- endif %} - {%- endif %} - {%- set reasoning_content = reasoning_content|trim %} - {%- if (loop.index0 > ns.last_query_index) or (keep_all_reasoning is defined and keep_all_reasoning) %} - {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content + '\n\n\n' + content }} - {%- else %} - {{- '<|im_start|>' + message.role + '\n' + content }} - {%- endif %} - {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} - {%- for tool_call in message.tool_calls %} - {%- if tool_call.function is defined %} - {%- set tool_call = tool_call.function %} - {%- endif %} - {%- if loop.first %} - {%- if content|trim %} - {{- '\n\n\n\n' }} - {%- else %} - {{- '\n\n' }} - {%- endif %} - {%- else %} - {{- '\n\n\n' }} - {%- endif %} - {%- if tool_call.arguments is defined %} - {%- for args_name, args_value in tool_call.arguments|items %} - {{- '\n' }} - {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} - {{- args_value }} - {{- '\n\n' }} - {%- endfor %} - {%- endif %} - {{- '\n' }} - {%- endfor %} - {%- endif %} - {{- '<|im_end|>\n' }} - {%- elif message.role == "tool" %} - {%- if loop.previtem and loop.previtem.role != "tool" %} - {{- '<|im_start|>user' }} - {%- endif %} - {{- '\n\n' }} - {{- content }} - {{- '\n' }} - {%- if not loop.last and loop.nextitem.role != "tool" %} - {{- '<|im_end|>\n' }} - {%- elif loop.last %} - {{- '<|im_end|>\n' }} - {%- endif %} - {%- else %} - {{- raise_exception('Unexpected message role.') }} - {%- endif %} -{%- endfor %} -{%- if add_generation_prompt %} - {{- '<|im_start|>assistant\n' }} - {%- if enable_thinking is defined and enable_thinking is false %} - {{- '\n\n\n\n' }} - {%- else %} - {{- '\n' }} - {%- endif %} -{%- endif %} diff --git a/inference/mlx/tests/test_e2e.py b/inference/mlx/tests/test_e2e.py deleted file mode 100644 index d054fbc..0000000 --- a/inference/mlx/tests/test_e2e.py +++ /dev/null @@ -1,229 +0,0 @@ -"""End-to-end (offline) checks on a tiny raw-HF checkpoint: loader, convert (q8), engine -prompt cache, OpenAI server (non-stream + SSE + tools), and output parsing.""" -import json -import os -import socket -import sys -import tempfile -import threading -import urllib.error -import urllib.request -from pathlib import Path - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -import mlx.core as mx - -import tiny_ckpt -from test_tiny import tiny_runtime -from trida_mlx import convert -from trida_mlx.decode import DecodeStats, SamplingParams -from trida_mlx.engine import Engine -from trida_mlx.model import TridaRuntime, load_model -from trida_mlx.parsing import StreamSplitter, parse_tool_calls, split_reasoning - -TMP = Path(tempfile.mkdtemp(prefix="trida_e2e_")) -CKPT, MASK_ID = tiny_ckpt.build(TMP / "raw") - - -def test_raw_hf_loader_matches_reference(): - model, cfg, _ = load_model(str(CKPT), dtype=mx.float32) - rt = TridaRuntime(model) - ref = tiny_runtime(seed=0) - p = [5, 17, 99, 3, 250, 8, 41] - a = rt.prefill(rt.make_cache(), p) - b = ref.prefill(ref.make_cache(), p) - # bf16 round trip of the weights -> small differences only - assert float(mx.max(mx.abs(a - b))) < 0.25 * float(mx.max(mx.abs(b))), (a[:5], b[:5]) - assert int(mx.argmax(a)) == int(mx.argmax(b)) - - -def test_convert_q8_and_selfspec_lossless_on_quantized(): - out = TMP / "q8" - convert.main(["--model", str(CKPT), "--out", str(out), "--bits", "8", "--group-size", "32"]) - eng = Engine(str(out)) - assert eng.mask_id == MASK_ID - p = eng.encode(eng.render([{"role": "user", "content": "hello world"}])) - sp = SamplingParams(temperature=0.0) - a = [t for c in eng.generate_ids(p, max_tokens=80, sampling=sp, mode="causal") for t in c] - eng.reset_cache() - b = [t for c in eng.generate_ids(p, max_tokens=80, sampling=sp, mode="self-spec") for t in c] - assert a == b - - -def test_prompt_cache_reuse(): - eng = Engine(str(CKPT)) - sp = SamplingParams(temperature=0.0) - msgs = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello world"}] - p1 = eng.encode(eng.render(msgs)) - out1 = [t for c in eng.generate_ids(p1, max_tokens=30, sampling=sp) for t in c] - msgs2 = msgs + [{"role": "assistant", "content": eng.tokenizer.decode(out1)}, - {"role": "user", "content": "and the weather?"}] - p2 = eng.encode(eng.render(msgs2)) - st = DecodeStats() - out2 = [t for c in eng.generate_ids(p2, max_tokens=30, sampling=sp, stats=st) for t in c] - assert st.reused_tokens > len(p1) // 2, st.as_dict() - # same result as a cold engine - eng2 = Engine(str(CKPT), prompt_cache=False) - ref2 = [t for c in eng2.generate_ids(p2, max_tokens=30, sampling=sp) for t in c] - assert out2 == ref2 - - -def test_parsing(): - c, calls = parse_tool_calls('ok {"name": "f", "arguments": {"x": 1}}}') - assert c == "ok" and calls == [{"name": "f", "arguments": {"x": 1}}] - _, calls = parse_tool_calls("\n\n\nSeoul\n\n\n") - assert calls == [{"name": "g", "arguments": {"city": "Seoul"}}] - assert split_reasoning("abc\n\nhi", True) == ("abc", "hi") - s = StreamSplitter(True, parse_tools=True) - text = 'let me think\n\nSure.{"name": "f", "arguments": {}}' - parts = [] - for ch in text: - parts += s.feed(ch) - rest, calls = s.finish() - parts += rest - reasoning = "".join(t for k, t in parts if k == "reasoning") - content = "".join(t for k, t in parts if k == "content") - assert reasoning == "let me think" and content == "Sure." and calls[0]["name"] == "f", (parts, calls) - - -def test_normalize_messages_against_real_trida_template(): - """Hermes-style histories must render with the real Trida chat template.""" - from transformers import AutoTokenizer - from trida_mlx.server import _normalize_messages - tok = AutoTokenizer.from_pretrained(str(CKPT)) - tok.chat_template = (Path(__file__).parent / "fixtures" / "trida_chat_template.jinja").read_text() - tools = [{"type": "function", "function": {"name": "terminal", "description": "run", - "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}}}] - msgs = [ - {"role": "system", "content": "You are Hermes."}, - {"role": "developer", "content": "Be terse."}, - {"role": "user", "content": [{"type": "text", "text": "list files"}]}, - {"role": "assistant", "content": None, "reasoning": "use terminal", - "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "terminal", "arguments": '{"command": "ls"}'}}]}, - {"role": "tool", "tool_call_id": "c1", "content": "a.py\nb.py"}, - {"role": "system", "content": "Context was compacted."}, - {"role": "user", "content": "thanks, now count them"}, - ] - norm = _normalize_messages(msgs) - assert [m["role"] for m in norm] == ["system", "user", "assistant", "tool", "user", "user"] - text = tok.apply_chat_template(norm, tools=tools, add_generation_prompt=True, tokenize=False, enable_thinking=True) - assert "You are Hermes.\n\nBe terse." in text - assert "\n\nls\n" in text - assert "[System note]\nContext was compacted." in text and text.endswith("\n") - - -def test_tool_param_coercion(): - tools = [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": { - "path": {"type": "string"}, "n": {"type": "integer"}, "flags": {"type": "array"}}}}}] - _, calls = parse_tool_calls("\n\n\n007\n\n" - "\n5\n\n\n[\"-a\"]\n\n" - "\n", tools) - assert calls == [{"name": "f", "arguments": {"path": "007", "n": 5, "flags": ["-a"]}}], calls - - -def test_multi_slot_cache_survives_side_requests(): - eng = Engine(str(CKPT), cache_slots=3) - sp = SamplingParams(temperature=0.0) - run = lambda msgs, st=None: [t for c in eng.generate_ids(eng.encode(eng.render(msgs)), max_tokens=12, - sampling=sp, stats=st) for t in c] - main = [{"role": "system", "content": "main agent " * 30}, {"role": "user", "content": "hello world"}] - out = run(main) - run([{"role": "user", "content": "write a title for: hello"}]) # side request (title) - run([{"role": "user", "content": "summarize: the weather in seoul"}]) # another side request - st = DecodeStats() - run(main + [{"role": "assistant", "content": eng.tokenizer.decode(out)}, - {"role": "user", "content": "and then?"}], st) - assert st.reused_tokens > 100, st.as_dict() - assert len(eng.slots) == 3 - - -def test_system_block_snapshot_reused_when_user_turn_changes(): - """Hermes rewrites the last user message on a truncation retry; the (long) system block - must still be reused.""" - eng = Engine(str(CKPT), cache_slots=1) - sp = SamplingParams(temperature=0.0) - sys_msg = {"role": "system", "content": "tool schemas " * 60} - p1 = eng.encode(eng.render([sys_msg, {"role": "user", "content": "list the files"}])) - for _ in eng.generate_ids(p1, max_tokens=5, sampling=sp): - pass - st = DecodeStats() - p2 = eng.encode(eng.render([sys_msg, {"role": "user", "content": "list the files [System: continue]"}])) - out = [t for c in eng.generate_ids(p2, max_tokens=5, sampling=sp, stats=st) for t in c] - n_sys = p1.index(eng.im_start, 1) # tokens up to the end of the system block - assert st.reused_tokens == n_sys, (st.reused_tokens, n_sys) - ref = Engine(str(CKPT), prompt_cache=False) - assert out == [t for c in ref.generate_ids(p2, max_tokens=5, sampling=sp) for t in c] - - -def _free_port(): - s = socket.socket(); s.bind(("127.0.0.1", 0)); p = s.getsockname()[1]; s.close(); return p - - -def test_server_roundtrip(): - from trida_mlx import server - port = _free_port() - t = threading.Thread(target=server.main, args=(["--model", str(CKPT), "--port", str(port), - "--max-tokens", "40", "--prefill-chunk", "16", - "--context-length", "2048"],), daemon=True) - t.start() - base = f"http://127.0.0.1:{port}/v1" - import time - for _ in range(200): - try: - urllib.request.urlopen(f"{base}/models", timeout=1); break - except Exception: - time.sleep(0.1) - tools = [{"type": "function", "function": {"name": "now", "description": "time", - "parameters": {"type": "object", "properties": {}}}}] - body = {"messages": [{"role": "user", "content": "hello"}], "tools": tools, "temperature": 0} - r = json.loads(urllib.request.urlopen(urllib.request.Request( - f"{base}/chat/completions", data=json.dumps(body).encode(), - headers={"Content-Type": "application/json"})).read()) - assert r["choices"][0]["message"]["role"] == "assistant" - assert r["usage"]["trida"]["new_tokens"] > 0 - body["stream"] = True - lines = urllib.request.urlopen(urllib.request.Request( - f"{base}/chat/completions", data=json.dumps(body).encode(), - headers={"Content-Type": "application/json"})).read().decode().splitlines() - datas = [l[6:] for l in lines if l.startswith("data: ")] - assert datas[-1] == "[DONE]" - usage = json.loads(datas[-2])["usage"] - assert usage["prompt_tokens_details"]["cached_tokens"] > 0 # second identical request hits the cache - assert any(l.startswith(": prefill") for l in lines) or usage["prompt_tokens_details"]["cached_tokens"] > 0 - r = json.loads(urllib.request.urlopen(urllib.request.Request( - f"{base}/completions", data=json.dumps({"prompt": "hello", "max_tokens": 10}).encode(), - headers={"Content-Type": "application/json"})).read()) - assert r["usage"]["completion_tokens"] <= 10 - models = json.loads(urllib.request.urlopen(f"{base}/models").read()) - assert models["data"][0]["context_length"] == 2048 and models["data"][0]["id"] == "trida2.0-4b" - # context overflow -> OpenAI-style 400 that harnesses recognise (and compress) - big = {"messages": [{"role": "user", "content": "hello world " * 3000}], "stream": True} - try: - urllib.request.urlopen(urllib.request.Request(f"{base}/chat/completions", data=json.dumps(big).encode(), - headers={"Content-Type": "application/json"})) - assert False, "expected 400" - except urllib.error.HTTPError as e: - err = json.loads(e.read())["error"] - assert e.code == 400 and err["code"] == "context_length_exceeded" and "maximum context length" in err["message"] - # keep-alive comments during a long (uncached) prefill - body2 = {"messages": [{"role": "user", "content": "the weather in seoul is sunny " * 40}], "stream": True, - "max_tokens": 4} - lines = urllib.request.urlopen(urllib.request.Request( - f"{base}/chat/completions", data=json.dumps(body2).encode(), - headers={"Content-Type": "application/json"})).read().decode().splitlines() - assert any(l.startswith(": prefill") for l in lines), lines[:5] - - -if __name__ == "__main__": - fails = 0 - for name, fn in list(globals().items()): - if name.startswith("test_"): - try: - fn(); print(f" ok {name}") - except Exception as e: # noqa: BLE001 - import traceback; traceback.print_exc() - fails += 1; print(f" FAIL {name}: {type(e).__name__}: {e}") - sys.stdout.flush() - os._exit(1 if fails else 0) diff --git a/inference/mlx/tests/test_tiny.py b/inference/mlx/tests/test_tiny.py deleted file mode 100644 index 7c9f01e..0000000 --- a/inference/mlx/tests/test_tiny.py +++ /dev/null @@ -1,297 +0,0 @@ -"""Correctness tests on a tiny random Qwen3.5-hybrid model (CPU is fine; no weights needed). - - python -m pytest inference/mlx/tests -q # or: python inference/mlx/tests/test_tiny.py - -What is pinned: - * chunked prefill == one-shot prefill - * canvas clean rows (0..N-1) reproduce the AR logits exactly (the lossless property) - * canvas MASK rows read the *block-end* GDN state (causal_mode=2 semantics) - * commit(adv) leaves the cache identical to having AR-decoded those adv tokens - * greedy self-spec output == greedy AR output, token for token, over long generations - * a snapshot/restore round trip (prompt cache) reproduces a fresh prefill -""" -import os -import sys - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -import mlx.core as mx -import mlx.nn as nn -import numpy as np - -from trida_mlx.decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate -from trida_mlx.model import GDNState, TridaRuntime, build_model -from mlx_lm.models.gated_delta import compute_g - -MASK = 299 -TINY = dict( - model_type="qwen3_5_text", hidden_size=64, intermediate_size=128, num_hidden_layers=8, - num_attention_heads=4, num_key_value_heads=2, head_dim=16, rms_norm_eps=1e-6, vocab_size=300, - linear_num_value_heads=4, linear_num_key_heads=2, linear_key_head_dim=32, linear_value_head_dim=16, - linear_conv_kernel_dim=4, tie_word_embeddings=True, full_attention_interval=4, - rope_parameters={"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.25, - "mrope_section": [1, 1, 0]}, -) - - -def tiny_runtime(seed=0, sharp=4.0): - mx.random.seed(seed) - model = build_model(TINY) - # sharpen the output distribution so greedy decoding has structure (and fewer near-ties) - emb = model.language_model.model.embed_tokens - emb.weight = emb.weight * sharp - mx.eval(model.parameters()) - return TridaRuntime(model, prefill_chunk=512) - - -def rand_prompt(n, seed=1): - rng = np.random.default_rng(seed) - return rng.integers(0, 290, size=n).tolist() - - -def close(a, b, tol=1e-4): - return float(mx.max(mx.abs(a - b)).item()) < tol - - -def test_chunked_prefill(): - rt = tiny_runtime() - p = rand_prompt(37) - c1 = rt.make_cache(); l1 = rt.prefill(c1, p) - rt2 = tiny_runtime(); rt2.prefill_chunk = 5 - c2 = rt2.make_cache(); l2 = rt2.prefill(c2, p) - assert close(l1, l2) - - -def _ar_logits(rt, prompt, cont): - c = rt.make_cache() - out = [rt.prefill(c, prompt)] - for t in cont: - out.append(rt.ar_step(c, t)) - return out, c - - -def test_canvas_clean_rows_are_ar_logits(): - rt = tiny_runtime(); n = 4 - prompt = rand_prompt(20) - cont = rand_prompt(n, seed=7) # p, s0, s1, s2 - ar, _ = _ar_logits(rt, prompt, cont) - c = rt.make_cache(); rt.prefill(c, prompt) - L = rt.canvas(c, cont + [MASK] * (n - 1), n) - for i in range(n): # row i == AR logits after cont[:i+1] - assert close(L[i], ar[i + 1]), f"row {i}" - - -def test_commit_matches_ar_state(): - n = 4 - for adv in range(1, n + 1): - rt = tiny_runtime() - prompt = rand_prompt(15) - cont = rand_prompt(n, seed=3) - c = rt.make_cache(); rt.prefill(c, prompt) - rt.canvas(c, cont + [MASK] * (n - 1), n) - rt.commit(c, adv) - nxt = 42 - got = rt.ar_step(c, nxt) - ar, _ = _ar_logits(rt, prompt, cont[:adv] + [nxt]) - assert close(got, ar[-1]), f"adv={adv}" - assert c.tokens == prompt + cont[:adv] + [nxt] - - -def test_mask_rows_read_block_end_state(): - """Independent re-derivation of one GDN layer's canvas output with a per-token loop.""" - rt = tiny_runtime(); n = 4 - layer = next(l for l in rt.backbone.layers if l.is_linear) - lin = layer.linear_attn - x = mx.random.normal((1, 2 * n - 1, TINY["hidden_size"])) - st = GDNState(conv=mx.random.normal((1, 3, lin.conv_dim)), - ssm=mx.random.normal((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim)) * 0.1) - out, _, _ = rt._gdn(lin, x, st, n) - # manual: conv -> q,k,v -> step recurrence, record states - qkv = lin.in_proj_qkv(x); z = lin.in_proj_z(x).reshape(1, 2 * n - 1, lin.num_v_heads, lin.head_v_dim) - ci = mx.concatenate([st.conv, qkv], axis=1) - co = nn.silu(lin.conv1d(ci)) - q, k, v = mx.split(co, [lin.key_dim, 2 * lin.key_dim], -1) - T = 2 * n - 1 - q = q.reshape(1, T, lin.num_k_heads, -1); k = k.reshape(1, T, lin.num_k_heads, -1) - v = v.reshape(1, T, lin.num_v_heads, -1) - s = k.shape[-1] ** -0.5 - q = s * s * mx.fast.rms_norm(q, None, 1e-6); k = s * mx.fast.rms_norm(k, None, 1e-6) - rep = lin.num_v_heads // lin.num_k_heads - q = mx.repeat(q, rep, 2); k = mx.repeat(k, rep, 2) - beta = mx.sigmoid(lin.in_proj_b(x)); g = compute_g(lin.A_log, lin.in_proj_a(x), lin.dt_bias) - S = st.ssm; ys = [] - for t in range(T): - S = S * g[:, t][..., None, None] - kv = (S * k[:, t][:, :, None, :]).sum(-1) - S = S + k[:, t][:, :, None, :] * ((v[:, t] - kv) * beta[:, t][..., None])[..., None] - ys.append((S * q[:, t][:, :, None, :]).sum(-1)) - S_end = S - for t in range(n, T): - ys[t] = (S_end * q[:, t][:, :, None, :]).sum(-1) - y = mx.stack(ys, axis=1) - ref = lin.out_proj(lin.norm(y, z).reshape(1, T, -1)) - assert close(out, ref, 1e-4) - - -def test_fused_gdn_matches_unfused(): - """Metal-only: the fused canvas GDN kernel == two mlx-lm kernel calls + block-end readout.""" - from trida_mlx.kernels import fused_available - if not fused_available(): - print(" (skipped: no Metal)") - return - rt = tiny_runtime(); n = 4 - lin = next(l for l in rt.backbone.layers if l.is_linear).linear_attn - x = mx.random.normal((1, 2 * n - 1, TINY["hidden_size"])) - st = GDNState(conv=mx.random.normal((1, 3, lin.conv_dim)), - ssm=mx.random.normal((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim)) * 0.1) - rt.fused_gdn = True - o1, _, r1 = rt._gdn(lin, x, st, n) - rt.fused_gdn = False - o2, _, r2 = rt._gdn(lin, x, st, n) - assert close(o1, o2, 1e-5) and close(r1.s_clean, r2.s_clean, 1e-6) - # clean rows must be bit-identical to the unfused (= AR) kernel - assert bool(mx.array_equal(o1[:, :n], o2[:, :n]).item()) - - -def _collect(gen): - out = [] - for chunk in gen: - out.extend(chunk) - return out - - -def test_greedy_selfspec_is_lossless(): - sp = SamplingParams(temperature=0.0) - for seed in range(3): - rt = tiny_runtime(seed=seed) - prompt = rand_prompt(25, seed=seed + 10) - ar = _collect(ar_generate(rt, rt.make_cache(), prompt, max_new_tokens=120, eos_ids=set(), sp=sp)) - from trida_mlx.decode import DecodeStats - st = DecodeStats() - ss = _collect(selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=120, eos_ids=set(), - sp=sp, mask_id=MASK, n=4, stats=st)) - assert ss == ar, f"seed {seed}: first diff at {next(i for i,(a,b) in enumerate(zip(ss,ar)) if a!=b)}" - assert st.forwards <= len(ar) - - -def test_greedy_lossless_with_oracle_drafts(): - """Random models rarely accept drafts, so also drive the accept path: patch the canvas - so MASK rows 'predict' the true AR continuation and check output + acceptance.""" - sp = SamplingParams(temperature=0.0) - rt = tiny_runtime(seed=5) - prompt = rand_prompt(18, seed=4) - ar = _collect(ar_generate(rt, rt.make_cache(), prompt, max_new_tokens=60, eos_ids=set(), sp=sp)) - full = prompt + ar - orig = rt.canvas - - # replace the MASK draft rows with a one-hot on the true AR continuation - def patched(cache, toks, n, rows=None): - base = cache.length - L = orig(cache, toks, n, rows=rows) - rows = [] - V = L.shape[-1] - for i in range(L.shape[0]): - want = base + i + 1 - if i >= 1 and want < len(full) and (toks[i] == MASK): - rows.append(mx.where(mx.arange(V) == full[want], 100.0, 0.0)) - else: - rows.append(L[i]) - return mx.stack(rows) - - rt.canvas = patched - from trida_mlx.decode import DecodeStats - st = DecodeStats() - ss = _collect(selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=60, eos_ids=set(), - sp=sp, mask_id=MASK, n=4, stats=st)) - assert ss == ar - assert st.accept_hist[-1] > 0 and st.tokens_per_forward > 2.5, st.as_dict() - - -def test_snapshot_restore(): - rt = tiny_runtime() - p = rand_prompt(30) - c = rt.make_cache(); rt.prefill(c, p[:20]) - snap = c.snapshot() - rt.prefill(c, p[20:]) - for t in rand_prompt(9, seed=9): - rt.ar_step(c, t) - c.restore(snap) - l1 = rt.prefill(c, p[20:]) - l2 = rt.prefill(rt.make_cache(), p) - assert close(l1, l2) - - -def test_sampling_runs(): - sp = SamplingParams(temperature=1.0, top_k=20, top_p=0.9, seed=0) - rt = tiny_runtime() - out = _collect(selfspec_generate(rt, rt.make_cache(), rand_prompt(12), max_new_tokens=40, - eos_ids=set(), sp=sp, mask_id=MASK, n=4)) - assert len(out) == 40 - - -def test_truncated_probs_matches_reference(): - from trida_mlx.decode import truncated_probs - - def ref(x, t, k, p): # numpy port of the SGLang reference _sampling_probs - x = x / t - kth = np.sort(x, axis=-1)[:, -k][:, None] - x = np.where(x < kth, -np.inf, x) - order = np.argsort(-x, axis=-1, kind="stable") - sx = np.take_along_axis(x, order, -1) - e = np.exp(sx - sx.max(-1, keepdims=True)); pr = e / e.sum(-1, keepdims=True) - sx = np.where(np.cumsum(pr, -1) - pr >= p, -np.inf, sx) - out = np.full_like(x, -np.inf); np.put_along_axis(out, order, sx, -1) - e = np.exp(out - out.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True) - - rng = np.random.default_rng(0) - x = rng.normal(size=(7, 1000)).astype(np.float32) * 3 - for t, k, p in [(1.0, 50, 0.95), (0.7, 20, 0.8), (1.0, 5, 1.0)]: - got = np.array(truncated_probs(mx.array(x), SamplingParams(temperature=t, top_k=k, top_p=p))) - assert np.allclose(got, ref(x, t, k, p), atol=1e-5), (t, k, p) - - -def test_speculative_sampling_preserves_distribution(): - """Sampled self-spec must reproduce the AR sampling distribution. Random models almost - never accept drafts, so the MASK rows are patched to reuse a neighbouring row's logits: - any draft q is valid for speculative sampling, and this one gets accepted often.""" - from collections import Counter - rt = tiny_runtime(seed=2, sharp=1.0) - prompt = rand_prompt(10, seed=2) - sp = SamplingParams(temperature=1.0, top_k=4, top_p=0.9) - orig = rt.canvas - - def patched(cache, toks, n, rows=None): - L = orig(cache, toks, n, rows=rows) - return mx.stack([L[i] if (i == 0 or toks[i] != MASK) else L[i - 1] for i in range(L.shape[0])]) - - def hist(fn, reps, seed0, **kw): - c, st_all = Counter(), DecodeStats() - st_all.accept_hist = [0] * 4 - for i in range(reps): - mx.random.seed(seed0 + i) - st = DecodeStats() - out = _collect(fn(rt, rt.make_cache(), prompt, max_new_tokens=4, eos_ids=set(), sp=sp, stats=st, **kw)) - c[out[3]] += 1 - if st.accept_hist: - st_all.accept_hist = [x + y for x, y in zip(st_all.accept_hist, st.accept_hist)] - return c, st_all - - R = 600 - a, _ = hist(ar_generate, R, 5000) - a2, _ = hist(ar_generate, R, 90000) - rt.canvas = patched - b, st = hist(selfspec_generate, R, 7000, mask_id=MASK, n=4) - tv = lambda x, y: 0.5 * sum(abs(x[k] - y[k]) for k in set(x) | set(y)) / R - assert sum(st.accept_hist[1:]) > R // 4, st.accept_hist # the accept path really ran - assert tv(a, b) < max(0.08, 2.0 * tv(a, a2)), (tv(a, b), tv(a, a2)) - - -if __name__ == "__main__": - fails = 0 - for name, fn in sorted(globals().items()): - if name.startswith("test_"): - try: - fn(); print(f" ok {name}") - except Exception as e: # noqa: BLE001 - fails += 1; print(f" FAIL {name}: {type(e).__name__}: {e}") - sys.exit(1 if fails else 0) diff --git a/inference/mlx/tests/tiny_ckpt.py b/inference/mlx/tests/tiny_ckpt.py deleted file mode 100644 index 31474c0..0000000 --- a/inference/mlx/tests/tiny_ckpt.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Build a tiny *raw-HF-format* Trida-like checkpoint (random weights, byte-level BPE -tokenizer, Qwen-style chat template with tools + thinking) for offline end-to-end tests.""" -import json -import os -import sys -from pathlib import Path - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -import mlx.core as mx -from mlx.utils import tree_flatten - -from test_tiny import TINY, tiny_runtime # noqa: E402 - -SPECIALS = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "", "", "", - "", "<|mask|>"] - -CHAT_TEMPLATE = r"""{%- if tools %}<|im_start|>system -{%- if messages[0].role == 'system' %} -{{ messages[0].content }} -{%- endif %} -# Tools - -{%- for t in tools %} -{{ t | tojson }} -{%- endfor %} - -Call a tool with {"name": ..., "arguments": {...}}<|im_end|> -{% elif messages[0].role == 'system' %}<|im_start|>system -{{ messages[0].content }}<|im_end|> -{% endif %} -{%- for m in messages %} -{%- if m.role == 'user' %}<|im_start|>user -{{ m.content }}<|im_end|> -{% elif m.role == 'assistant' %}<|im_start|>assistant -{{ m.content }} -{%- if m.tool_calls %}{% for tc in m.tool_calls %}{"name": "{{ tc.function.name }}", "arguments": {{ tc.function.arguments | tojson }}}{% endfor %}{% endif %}<|im_end|> -{% elif m.role == 'tool' %}<|im_start|>user -{{ m.content }}<|im_end|> -{% endif %} -{%- endfor %} -{%- if add_generation_prompt %}<|im_start|>assistant -{%- if enable_thinking is defined and enable_thinking is false %} - - - - -{% else %} - -{% endif %} -{%- endif %}""" - - -def build(out: Path): - from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers - from transformers import PreTrainedTokenizerFast - - out.mkdir(parents=True, exist_ok=True) - tok = Tokenizer(models.BPE()) - tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) - tok.decoder = decoders.ByteLevel() - corpus = ["hello world the weather in seoul is sunny", "def fib(n): return n", "tool call json name", - "안녕하세요 온디바이스"] * 20 - trainer = trainers.BpeTrainer(vocab_size=TINY["vocab_size"] - len(SPECIALS), special_tokens=[], - initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False) - tok.train_from_iterator(corpus, trainer) - tok.add_special_tokens(SPECIALS) - assert tok.get_vocab_size() <= TINY["vocab_size"], tok.get_vocab_size() - fast = PreTrainedTokenizerFast(tokenizer_object=tok, eos_token="<|im_end|>", pad_token="<|endoftext|>") - fast.chat_template = CHAT_TEMPLATE - fast.save_pretrained(str(out)) - mask_id = fast.convert_tokens_to_ids("<|mask|>") - - rt = tiny_runtime(seed=0) - params = dict(tree_flatten(rt.model.parameters())) - hf = {} - norm_sfx = (".input_layernorm.weight", ".post_attention_layernorm.weight", "model.norm.weight", - ".q_norm.weight", ".k_norm.weight") - for k, v in params.items(): - k2 = k.replace("language_model.model.", "model.").replace("language_model.", "") - if "conv1d.weight" in k2: - v = v.moveaxis(1, 2) # MLX (C, K, 1) -> HF (C, 1, K) - if any(k2.endswith(s) for s in norm_sfx): - v = v - 1.0 # HF Qwen3.5 stores zero-centred RMSNorm weights - hf[k2] = v.astype(mx.bfloat16) if not k2.endswith("A_log") else v - hf["lm_head.weight"] = hf["model.embed_tokens.weight"] # real ckpt ships the tied head too - mx.save_safetensors(str(out / "model.safetensors"), hf) - cfg = dict(TINY, architectures=["Qwen3_5ForCausalLM"]) - (out / "config.json").write_text(json.dumps(cfg, indent=2)) - (out / "block_diffusion.json").write_text(json.dumps({"bd_size": 4, "mask_id": mask_id})) - (out / "generation_config.json").write_text(json.dumps({"eos_token_id": [fast.eos_token_id]})) - return out, mask_id - - -if __name__ == "__main__": - print(build(Path(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tiny_trida"))) diff --git a/inference/mlx/uv.lock b/inference/mlx/uv.lock index fdcb1ff..89282bf 100644 --- a/inference/mlx/uv.lock +++ b/inference/mlx/uv.lock @@ -3,10 +3,10 @@ revision = 3 requires-python = ">=3.10, <3.14" resolution-markers = [ "python_full_version >= '3.12' and platform_machine == 'arm64' and sys_platform == 'darwin'", - "python_full_version == '3.11.*' and platform_machine == 'arm64' and sys_platform == 'darwin'", - "python_full_version < '3.11' and platform_machine == 'arm64' and sys_platform == 'darwin'", "python_full_version >= '3.12' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'arm64' and sys_platform == 'darwin'", "python_full_version == '3.11.*' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'arm64' and sys_platform == 'darwin'", "python_full_version < '3.11' and sys_platform == 'linux'", ] supported-markers = [ @@ -165,15 +165,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/58/a2/bb081bab032533a855d44de1d56f8e8426114ff1ba5d1f07a438a0a654f8/idna-3.20-py3-none-any.whl", hash = "sha256:ab7ae7122974553370f0bdb919e1a960b2cd1bc1ef0276416d896db81c14582c", size = 69583, upload-time = "2026-09-17T14:11:03.168Z" }, ] -[[package]] -name = "iniconfig" -version = "2.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, -] - [[package]] name = "jinja2" version = "3.1.6" @@ -445,15 +436,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, ] -[[package]] -name = "pluggy" -version = "1.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, -] - [[package]] name = "protobuf" version = "7.36.2" @@ -476,23 +458,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, ] -[[package]] -name = "pytest" -version = "9.1.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, - { name = "iniconfig" }, - { name = "packaging" }, - { name = "pluggy" }, - { name = "pygments" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, -] - [[package]] name = "pyyaml" version = "6.0.3" @@ -742,11 +707,6 @@ dependencies = [ { name = "transformers" }, ] -[package.dev-dependencies] -dev = [ - { name = "pytest" }, -] - [package.metadata] requires-dist = [ { name = "huggingface-hub" }, @@ -756,9 +716,6 @@ requires-dist = [ { name = "transformers", specifier = ">=4.57" }, ] -[package.metadata.requires-dev] -dev = [{ name = "pytest", specifier = ">=8" }] - [[package]] name = "typer" version = "0.27.2" From b56cd6ad47aa07676a55e3168b66a458bdcc0542 Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Tue, 6 Oct 2026 01:09:03 +0000 Subject: [PATCH 5/9] feat(inference/mlx): image input for vision-language checkpoints (Trida-2.0-4B-1006) - trida_mlx/vision.py: Qwen3.5 vision encoder + projector (adapted from mlx-vlm, MIT), Qwen2-VL preprocessing (bit-identical to the HF PIL processor), interleaved multimodal RoPE - runtime: image tokens are per-image hash keys, so the prompt cache never confuses two images and follow-up turns don't re-encode them; prompts with images prefill with MRoPE, SeqCache.pos tracks the next RoPE position, decode and the self-spec canvas stay plain RoPE - loader/convert: load the vision tower when the checkpoint has one; keep it bf16 when quantizing; carry processor configs and transformers_version - server: OpenAI image_url parts (data URL / http / path), 400 on bad images, supports_vision on /v1/models, --max-image-pixels - verify --image; self-spec default stays N=4 (on M3 Pro q8 N=8 is 2.72 vs 2.40 tok/fwd but 43 vs 52 tok/s); README block-size table, NOTICE / COMPLIANCE for the mlx-vlm-derived encoder, pillow dependency Checked: full multimodal prefill matches mlx-vlm's Qwen3.5 to 4e-6 (fp32, tiny model); on M3 Pro the 1006 q8 model describes a screenshot correctly and greedy self-spec matches AR with a 975-token image prompt. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- CHANGELOG.md | 3 + COMPLIANCE.md | 1 + NOTICE | 8 + inference/mlx/README.md | 64 +++++- inference/mlx/pyproject.toml | 1 + inference/mlx/requirements.txt | 1 + inference/mlx/trida_mlx/bench.py | 2 +- inference/mlx/trida_mlx/convert.py | 20 +- inference/mlx/trida_mlx/engine.py | 63 +++++- inference/mlx/trida_mlx/model.py | 112 ++++++++-- inference/mlx/trida_mlx/profile.py | 8 +- inference/mlx/trida_mlx/server.py | 93 ++++++-- inference/mlx/trida_mlx/verify.py | 21 +- inference/mlx/trida_mlx/vision.py | 336 +++++++++++++++++++++++++++++ inference/mlx/uv.lock | 36 ++++ 15 files changed, 707 insertions(+), 62 deletions(-) create mode 100644 inference/mlx/trida_mlx/vision.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 016345b..32aa065 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `inference/mlx/`: on-device (Apple Silicon, MLX) serving of Trida2.0-4B with lossless self-speculative decoding (`bd_bidir_shift` b7/g4 semantics), an OpenAI-compatible server with tool calls and a prompt cache, a minimal agent loop, and MLX quantization. +- `inference/mlx/`: image input for vision-language checkpoints (Trida-2.0-4B-1006): Qwen3.5 vision + encoder + projector, HF-identical preprocessing, multimodal RoPE, image-aware prompt cache, + OpenAI `image_url` parts in the server. ## [0.1.0] — 2026-09-30 diff --git a/COMPLIANCE.md b/COMPLIANCE.md index ae12c2c..f6cd428 100644 --- a/COMPLIANCE.md +++ b/COMPLIANCE.md @@ -11,6 +11,7 @@ model-card license. | Component | Path | Upstream | License | |---|---|---|---| | IFEval scorer | `inference/ifeval_lib/` | [google-research · instruction_following_eval](https://github.com/google-research/google-research/tree/master/instruction_following_eval) | **Apache-2.0** | +| mlx-vlm Qwen3-VL vision encoder (adapted) | `inference/mlx/trida_mlx/vision.py` | [Blaizzy/mlx-vlm](https://github.com/Blaizzy/mlx-vlm) `mlx_vlm/models/qwen3_vl/vision.py` | **MIT** | | MLX gated-delta Metal kernel (adapted) | `inference/mlx/trida_mlx/kernels.py` | [ml-explore/mlx-lm](https://github.com/ml-explore/mlx-lm) `mlx_lm/models/gated_delta.py` | **MIT** | These are the only third-party source code redistributed in the repository (see diff --git a/NOTICE b/NOTICE index fea11f5..9f5f6dc 100644 --- a/NOTICE +++ b/NOTICE @@ -39,6 +39,14 @@ MLX gated-delta Metal kernel (mlx-lm) mlx-lm's gated_delta Metal kernel source (same thread layout and recurrence; adds the block-end readout for MASK rows). See its header. +mlx-vlm Qwen3-VL / Qwen3.5 vision encoder + Copyright © 2025 Prince Canuma. + License: MIT + Source: https://github.com/Blaizzy/mlx-vlm (mlx_vlm/models/qwen3_vl/vision.py) + Derived: inference/mlx/trida_mlx/vision.py — the vision tower (patch embed, interpolated + position embedding, 2-D rotary blocks, patch merger) is adapted from mlx-vlm so + parameter names match its checkpoints. See its header. + The components above are the third-party source code redistributed in this repository. They do not share a single license: see COMPLIANCE.md section 1. diff --git a/inference/mlx/README.md b/inference/mlx/README.md index 9d9000f..d4e48bf 100644 --- a/inference/mlx/README.md +++ b/inference/mlx/README.md @@ -1,14 +1,17 @@ # inference/mlx — on-device self-speculative decoding (Apple Silicon) -Runs **[`trillionlabs/Trida2.0-4B`](https://huggingface.co/trillionlabs/Trida2.0-4B)** on a Mac with -**MLX**, in the same three pieces as the GPU stacks next door — but in-process, one stream, no CUDA: +Runs **[`trillionlabs/Trida2.0-4B`](https://huggingface.co/trillionlabs/Trida2.0-4B)** and the +vision-language **[`trillionlabs/Trida-2.0-4B-1006`](https://huggingface.co/trillionlabs/Trida-2.0-4B-1006)** +on a Mac with **MLX**, in the same three pieces as the GPU stacks next door — but in-process, one +stream, no CUDA: | | file | what | |---|---|---| | decoder | `trida_mlx/model.py`, `trida_mlx/decode.py` | `causal` (AR baseline) and **`self-spec`** (diffusion draft + exact AR verify, lossless) | | server | `trida_mlx/server.py` | OpenAI-compatible `/v1/chat/completions` (+ `/v1/completions`), SSE streaming, **tool calls**, reasoning split, prompt cache | | agent | `trida_mlx/agent.py` | a tiny tool-using agent loop against the server (files, search, calculator; opt-in shell/write) | -| tools | `convert.py`, `verify.py`, `bench.py` | quantize to MLX, first-contact correctness check, causal-vs-self-spec benchmark | +| vision | `trida_mlx/vision.py` | Qwen3.5 vision encoder + projector, image preprocessing, multimodal RoPE (checkpoints with a `vision_config`) | +| tools | `convert.py`, `verify.py`, `bench.py`, `profile.py` | quantize to MLX, first-contact correctness check, causal-vs-self-spec benchmark, step profile | Everything is built on `mlx_lm.models.qwen3_5` (weights, quantized layers, RoPE and the Metal gated-delta kernel come from mlx-lm unchanged); this package only adds the self-spec forward, @@ -49,15 +52,50 @@ r = c.chat.completions.create(model="trida", messages=[{"role": "user", "content print(r.choices[0].message.content, r.usage) # usage.trida has tok/s, tok/forward, accept histogram ``` -Server flags: `--mode self-spec|causal`, `--gen-block N` (canvas `2N-1`, default 4 → the shipped -`b7/g4` config), `--temperature/--top-k/--top-p` defaults (1.0 / 50 / 0.95, the canonical Trida -self-spec setting), `--no-think`, `--max-tokens`, `--no-prompt-cache`. Per request: `trida_mode`, -`chat_template_kwargs`, `tools`, `stop`, `seed`, `stream`. +Server flags: `--mode self-spec|causal`, `--gen-block N` (canvas `2N-1`, default 4 → `b7_g4`; see +"Block size" below), `--temperature/--top-k/--top-p` defaults +(1.0 / 50 / 0.95, the canonical Trida self-spec setting), `--no-think`, `--max-tokens`, +`--no-prompt-cache`, `--max-image-pixels`. Per request: `trida_mode`, `chat_template_kwargs`, +`tools`, `stop`, `seed`, `stream`. + +## Block size + +Both checkpoints were trained with `bd_size` 8, and the 1006 model card's GPU self-spec config is +`b15_g8`. On a Mac the canvas is not free: M3 Pro 18 GB, q8 Trida-2.0-4B-1006, greedy, 512 tokens: + +| N (canvas) | causal tok/s | self-spec tok/s | tok/forward | step vs AR step | speedup | +|---|---:|---:|---:|---:|---:| +| 4 (7) | 26.2 | **52.4** | 2.40 | 1.18× | **2.00×** | +| 8 (15) | 25.9 | 43.0 | 2.72 | 1.64× | 1.66× | + +N = 8 accepts more per step but each step costs more than it gains, so the default is N = 4. + +## Images (Trida-2.0-4B-1006) + +Checkpoints in the vision-language layout (`Qwen3_5ForConditionalGeneration`, with `vision_config`) +load their Qwen3.5 vision encoder + projector automatically; `convert.py` quantizes the language +model and keeps the vision tower in bf16. Images go through the server as OpenAI content parts +(`{"type": "image_url", "image_url": {"url": "data:image/png;base64,..." | "https://..." | "/path"}}`): + +- preprocessing = Qwen2-VL smart resize to multiples of 32 px, normalize, 16×16×2 patches + (bit-identical to the HF `Qwen2VLImageProcessor` PIL path); `--max-image-pixels` (default + 1024×1024 → ≤ 1024 image tokens) caps the token cost on device; +- the template's `<|image_pad|>` is expanded to one token per merged 2×2 patch, image rows get the + vision features, and the prompt uses interleaved multimodal RoPE (t/h/w positions); text after an + image continues from `max(t, h, w)`, so decoding and the self-spec canvas stay plain RoPE; +- the prompt cache keys image tokens by a hash of the image, so a different image never reuses a + cache, and a follow-up turn about the same image does not re-encode it; +- checked offline against mlx-vlm's Qwen3.5 implementation: same positions, last-token logits + within 4e-6 in float32 on a random tiny model. + +`verify.py --image photo.jpg` runs the lossless/speed checks on an image prompt. In Hermes, set +`hermes config set model.supports_vision true` so images are sent natively. ## How self-spec maps onto MLX -Reference semantics are the SGLang `HybridDiffusionSelfSpec` (`bd_bidir_shift`, block 7, gen 4, -`causal_prefill`, `draft_mode: strict_truncated`) — the same algorithm the vLLM plugin ports. +Reference semantics are the SGLang `HybridDiffusionSelfSpec` (`bd_bidir_shift`, block `2N-1`, +gen `N`, `causal_prefill`, `draft_mode: strict_truncated`) — the same algorithm the vLLM plugin +ports. Shown for N = 4: ``` cold start : [t0, M, M, M, M, M, M ] clean = row 0, specs = rows 1..3 @@ -128,6 +166,12 @@ shows the reuse. attention / gated-delta and prints the break-even tokens-per-forward. Done: a fused canvas GDN Metal kernel (`trida_mlx/kernels.py`, one launch per layer; `--no-fused-gdn` to A/B): verify step 49.8 → 44.5 ms on M3 Pro q8, i.e. 1.13× an AR step (39.3 ms). LM-head row skipping on cold starts - is available (`bench --skip-rows`) but measured slower, so it is off. + is available (`bench --skip-rows`) but measured slower, so it is off. (N = 4.) +- **Tool calls.** The Trida-2.0-4B-1006 model card reports lower BFCL scores for self-spec than AR on + the SGLang path (unclosed ``). Here greedy self-spec reproduces AR token for token and + sampled self-spec keeps the AR distribution, so that gap should not appear — but measure it; and + `--mode causal` (or `"trida_mode": "causal"` per request) is the fallback. +- **Vision scope.** Still images only (no video), no deepstack features (the released checkpoints + have none); images inside system or assistant messages are dropped by the chat template contract. - The GGUF quants (`mradermacher/Trida2.0-4B-*GGUF`) run in llama.cpp as plain AR only; self-spec there would need the canvas mask and the block-end GDN readout in llama.cpp's qwen35 graph. diff --git a/inference/mlx/pyproject.toml b/inference/mlx/pyproject.toml index b8f761b..0c6f636 100644 --- a/inference/mlx/pyproject.toml +++ b/inference/mlx/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "mlx-lm>=0.31", # Qwen3.5 gated-delta model + Metal kernel we build on "transformers>=4.57", # tokenizer + chat template (tools / enable_thinking) "huggingface-hub", + "pillow", # image decoding / resizing for vision checkpoints ] [project.scripts] diff --git a/inference/mlx/requirements.txt b/inference/mlx/requirements.txt index 2830b61..d4606a8 100644 --- a/inference/mlx/requirements.txt +++ b/inference/mlx/requirements.txt @@ -3,3 +3,4 @@ mlx>=0.29 mlx-lm>=0.31 # provides the Qwen3.5 gated-delta model + Metal kernel we build on transformers>=4.57 # tokenizer + chat template (tools / enable_thinking) huggingface_hub +pillow # image input for vision checkpoints diff --git a/inference/mlx/trida_mlx/bench.py b/inference/mlx/trida_mlx/bench.py index 08d75f4..b0518cb 100644 --- a/inference/mlx/trida_mlx/bench.py +++ b/inference/mlx/trida_mlx/bench.py @@ -68,7 +68,7 @@ def main(argv=None): ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") ap.add_argument("--prompts", default=None) ap.add_argument("--modes", default="causal,self-spec") - ap.add_argument("--gen-block", type=int, default=4) + ap.add_argument("--gen-block", type=int, default=None, help="N (canvas 2N-1); default 4") ap.add_argument("--max-tokens", type=int, default=512) ap.add_argument("--temperature", type=float, default=0.0) ap.add_argument("--top-k", type=int, default=50) diff --git a/inference/mlx/trida_mlx/convert.py b/inference/mlx/trida_mlx/convert.py index abcaf2f..3f422c2 100644 --- a/inference/mlx/trida_mlx/convert.py +++ b/inference/mlx/trida_mlx/convert.py @@ -25,7 +25,11 @@ from .model import _text_config, load_model CARRY = ["tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "generation_config.json", - "block_diffusion.json", "special_tokens_map.json", "vocab.json", "merges.txt", "added_tokens.json"] + "block_diffusion.json", "special_tokens_map.json", "vocab.json", "merges.txt", "added_tokens.json", + "preprocessor_config.json", "processor_config.json", "video_preprocessor_config.json"] +# multimodal config keys carried next to the (flattened) text config +VISION_KEYS = ["vision_config", "image_token_id", "video_token_id", "vision_start_token_id", + "vision_end_token_id"] def main(argv=None): @@ -39,11 +43,18 @@ def main(argv=None): ap.add_argument("--shard-gb", type=float, default=4.0) a = ap.parse_args(argv) - model, cfg, src = load_model(a.model) + model, vision, cfg, src = load_model(a.model) text = _text_config(cfg) out_cfg = dict(text) out_cfg["model_type"] = "qwen3_5" out_cfg["architectures"] = ["Qwen3_5ForCausalLM"] + # transformers >= 5 treats a local tokenizer whose config.json lacks transformers_version as a + # possible old Mistral one and prints a (false) "incorrect regex pattern" warning + out_cfg.setdefault("transformers_version", cfg.get("transformers_version", text.get("transformers_version", "5.12.1"))) + if vision is not None: # the vision tower stays bf16 (small, and the most quant-sensitive part) + for k in VISION_KEYS: + if k in cfg: + out_cfg[k] = cfg[k] if a.bits: def pred(path, m): if not hasattr(m, "to_quantized"): @@ -57,6 +68,8 @@ def pred(path, m): out = Path(a.out) out.mkdir(parents=True, exist_ok=True) weights = dict(tree_flatten(model.parameters())) + if vision is not None: + weights.update({f"vision_tower.{k}": v for k, v in tree_flatten(vision.parameters())}) shards, cur, cur_bytes = [], {}, 0 limit = int(a.shard_gb * (1 << 30)) for k, v in weights.items(): @@ -78,7 +91,8 @@ def pred(path, m): shutil.copy2(src / name, out / name) gb = index["metadata"]["total_size"] / 1e9 print(f"wrote {out} ({gb:.2f} GB, {len(shards)} shard(s), " - f"{'bf16' if not a.bits else f'q{a.bits} g{a.group_size}'})") + f"{'bf16' if not a.bits else f'q{a.bits} g{a.group_size}'}" + f"{', + bf16 vision tower' if vision is not None else ''})") if __name__ == "__main__": diff --git a/inference/mlx/trida_mlx/engine.py b/inference/mlx/trida_mlx/engine.py index 89e338c..a5c89d2 100644 --- a/inference/mlx/trida_mlx/engine.py +++ b/inference/mlx/trida_mlx/engine.py @@ -21,6 +21,7 @@ from .decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate from .model import DEFAULT_MODEL, TridaRuntime, load_model, mask_id_of +from .vision import ImageProcessorConfig, image_key, load_image, preprocess class IncrementalDetokenizer: @@ -56,17 +57,26 @@ class Engine: harness interleaves its main loop with side requests (titles, summaries, memory review); with one slot every side request would evict the main conversation's cache.""" - def __init__(self, model: str = DEFAULT_MODEL, *, mode: str = "self-spec", gen_block: int = 4, + def __init__(self, model: str = DEFAULT_MODEL, *, mode: str = "self-spec", gen_block: Optional[int] = None, prefill_chunk: int = 512, prompt_cache: bool = True, fused_gdn: bool = True, - cache_slots: int = 3): + cache_slots: int = 3, max_image_pixels: int = 1024 * 1024, image_cache: int = 16): from transformers import AutoTokenizer t = time.perf_counter() - self.model, self.cfg, self.path = load_model(model) + self.model, self.vision, self.cfg, self.path = load_model(model) self.tokenizer = AutoTokenizer.from_pretrained(str(self.path)) self.mask_id = mask_id_of(self.path, self.cfg, self.tokenizer) - self.rt = TridaRuntime(self.model, prefill_chunk=prefill_chunk, fused_gdn=fused_gdn) - self.mode, self.n = mode, gen_block + self.image_token_id = int(self.cfg.get("image_token_id", 248056)) + self.image_pad = self.tokenizer.convert_tokens_to_ids("<|image_pad|>") + self.image_proc = ImageProcessorConfig.from_files(self.path, max_image_pixels) + self.image_cache = image_cache + self.rt = TridaRuntime(self.model, prefill_chunk=prefill_chunk, fused_gdn=fused_gdn, + vision=self.vision, image_token_id=self.image_token_id) + # N tokens per self-spec step (canvas 2N-1). Default 4: on an M3 Pro (q8, Trida-2.0-4B-1006) + # N=8 accepts more per step (2.72 vs 2.40 tok/fwd) but its 15-token canvas costs 1.64x an + # AR step vs 1.18x for N=4, so N=4 decodes faster (52.4 vs 43.0 tok/s). On GPUs, where + # extra canvas rows are nearly free, the checkpoint's bd_size (8, the b15_g8 config) wins. + self.mode, self.n = mode, gen_block or 4 self.model_id = model self.eos_ids = self._eos_ids() self.im_start = self.tokenizer.convert_tokens_to_ids("<|im_start|>") @@ -115,8 +125,47 @@ def render(self, messages: list, tools: Optional[list] = None, enable_thinking: enable_thinking=enable_thinking, **template_kwargs, ) - def encode(self, text: str) -> list: - return self.tokenizer.encode(text, add_special_tokens=False) + def encode(self, text: str, images: Optional[list] = None) -> list: + """Token ids; each ``<|image_pad|>`` the template emitted is expanded into the + image's key repeated once per merged patch (see ``add_image``).""" + ids = self.tokenizer.encode(text, add_special_tokens=False) + if not images: + return ids + if self.vision is None: + raise ValueError("this checkpoint has no vision encoder; images are not supported") + n_pads = sum(1 for t in ids if t == self.image_pad) + if n_pads != len(images): + raise ValueError(f"prompt has {n_pads} image placeholders but {len(images)} images") + out, it = [], iter(images) + for t in ids: + if t == self.image_pad: + key, n = next(it) + out.extend([key] * n) + else: + out.append(t) + return out + + @property + def supports_vision(self) -> bool: + return self.vision is not None + + def add_image(self, src) -> tuple: + """Decode + preprocess one image (any thread). Returns ``(key, n_tokens)``; the vision + features are computed lazily on the worker at prefill time and cached by content.""" + img, data = load_image(src) + key = image_key(data, self.image_proc) + ent = self.rt.images.get(key) + if ent is None: + pixels, grid = preprocess(img, self.image_proc) + ent = {"pixels": pixels, "grid": grid, "features": None} + self.rt.images[key] = ent + while len(self.rt.images) > self.image_cache: # drop the oldest entry + self.rt.images.pop(next(iter(self.rt.images))) + else: # refresh LRU position + self.rt.images[key] = self.rt.images.pop(key) + t, h, w = ent["grid"] + m = self.image_proc.merge_size + return key, t * (h // m) * (w // m) # -- prompt cache ----------------------------------------------------------- def _position_cache(self, prompt: list) -> None: diff --git a/inference/mlx/trida_mlx/model.py b/inference/mlx/trida_mlx/model.py index 17249ad..7294278 100644 --- a/inference/mlx/trida_mlx/model.py +++ b/inference/mlx/trida_mlx/model.py @@ -28,12 +28,14 @@ from typing import Optional import mlx.core as mx +import numpy as np import mlx.nn as nn from mlx_lm.models import qwen3_5 from mlx_lm.models.cache import KVCache from mlx_lm.models.gated_delta import compute_g, gated_delta_kernel, gated_delta_ops from .kernels import canvas_gdn_fused, fused_available +from .vision import VisionConfig, VisionTower, apply_mrope, is_vision_key, mrope_cos_sin, sanitize_vision_weights DEFAULT_MODEL = "trillionlabs/Trida2.0-4B" DEFAULT_MASK_ID = 248077 @@ -59,6 +61,8 @@ def resolve_model_path(path_or_id: str) -> Path: def _text_config(cfg: dict) -> dict: text = dict(cfg.get("text_config", cfg)) text.pop("architectures", None) + if "tie_word_embeddings" not in text and "tie_word_embeddings" in cfg: + text["tie_word_embeddings"] = cfg["tie_word_embeddings"] return text @@ -67,10 +71,11 @@ def build_model(text_cfg: dict) -> qwen3_5.Model: return qwen3_5.Model(args) -def load_model(path_or_id: str = DEFAULT_MODEL, *, dtype=mx.bfloat16): +def load_model(path_or_id: str = DEFAULT_MODEL, *, dtype=mx.bfloat16, load_vision: bool = True): """Load a raw HF Trida checkpoint or one written by ``convert.py`` (optionally quantized). - Returns ``(model, config_dict, path)``. + Returns ``(model, vision_tower_or_None, config_dict, path)``. The vision tower exists when + the checkpoint has a ``vision_config`` and vision weights (kept in bf16, never quantized). """ path = resolve_model_path(path_or_id) cfg = json.loads((path / "config.json").read_text()) @@ -80,6 +85,15 @@ def load_model(path_or_id: str = DEFAULT_MODEL, *, dtype=mx.bfloat16): weights.update(mx.load(str(f))) if not weights: raise FileNotFoundError(f"no *.safetensors in {path}") + vision_w = {k: weights.pop(k) for k in [k for k in weights if is_vision_key(k)]} + vision = None + if load_vision and vision_w and "vision_config" in cfg: + vision = VisionTower(VisionConfig.from_dict(cfg["vision_config"])) + vw = sanitize_vision_weights(vision_w) + vw = {k: (v.astype(dtype) if mx.issubdtype(v.dtype, mx.floating) else v) for k, v in vw.items()} + vision.load_weights(list(vw.items()), strict=True) + mx.eval(vision.parameters()) + vision.eval() quant = cfg.get("quantization") if quant: nn.quantize( @@ -109,7 +123,7 @@ def load_model(path_or_id: str = DEFAULT_MODEL, *, dtype=mx.bfloat16): model.load_weights(list(weights.items()), strict=True) mx.eval(model.parameters()) model.eval() - return model, cfg, path + return model, vision, cfg, path def mask_id_of(path: Path, cfg: dict, tokenizer=None) -> int: @@ -156,6 +170,9 @@ class SeqCache: tokens: list = field(default_factory=list) pending: Optional[list] = None # _CanvasGDN per layer after a canvas forward canvas_len: int = 0 + # next RoPE position. == len(tokens) for text; images advance it by max(t, h, w) of + # their merged grid instead of their token count (multimodal RoPE) + pos: int = 0 @property def length(self) -> int: @@ -168,10 +185,11 @@ def snapshot(self): return ( list(self.tokens), [l if isinstance(l, GDNState) else l.offset for l in self.layers], + self.pos, ) def restore(self, snap) -> None: - tokens, states = snap + tokens, states, pos = snap assert self.pending is None for i, s in enumerate(states): if isinstance(s, GDNState): @@ -182,6 +200,7 @@ def restore(self, snap) -> None: raise ValueError("cannot restore a KV snapshot longer than the cache") kv.trim(kv.offset - s) self.tokens = list(tokens) + self.pos = pos def nbytes(self) -> int: n = 0 @@ -202,8 +221,13 @@ def _gdr(q, k, v, g, beta, state): class TridaRuntime: - def __init__(self, model: qwen3_5.Model, prefill_chunk: int = 512, fused_gdn: bool = True): + def __init__(self, model: qwen3_5.Model, prefill_chunk: int = 512, fused_gdn: bool = True, + vision: Optional[VisionTower] = None, image_token_id: int = 248056): self.model = model + self.vision = vision + self.image_token_id = image_token_id + # negative token id (image key) -> {"pixels", "grid", "features"}; see Engine.add_image + self.images: dict = {} # one Metal launch per GDN layer for the canvas (GPU only; ops fallback elsewhere) self.fused_gdn = fused_gdn # cold start: LM head on rows 0..N-1 only. Off by default: measured slower on M3 Pro @@ -235,7 +259,7 @@ def make_cache(self) -> SeqCache: # -- layers ----------------------------------------------------------------- @staticmethod - def _attention(attn, x, kv: KVCache, mask): + def _attention(attn, x, kv: KVCache, mask, offset: int, mrope=None): B, L, _ = x.shape qo = attn.q_proj(x) queries, gate = mx.split(qo.reshape(B, L, attn.num_attention_heads, -1), 2, axis=-1) @@ -244,8 +268,11 @@ def _attention(attn, x, kv: KVCache, mask): queries = attn.q_norm(queries).transpose(0, 2, 1, 3) keys = attn.k_norm(keys.reshape(B, L, attn.num_key_value_heads, -1)).transpose(0, 2, 1, 3) values = values.reshape(B, L, attn.num_key_value_heads, -1).transpose(0, 2, 1, 3) - queries = attn.rope(queries, offset=kv.offset) - keys = attn.rope(keys, offset=kv.offset) + if mrope is not None: + queries, keys = apply_mrope(queries, *mrope), apply_mrope(keys, *mrope) + else: + queries = attn.rope(queries, offset=offset) + keys = attn.rope(keys, offset=offset) keys, values = kv.update_and_fetch(keys, values) out = mx.fast.scaled_dot_product_attention(queries, keys, values, scale=attn.scale, mask=mask) out = out.transpose(0, 2, 1, 3).reshape(B, L, -1) @@ -307,8 +334,9 @@ def _lm_head(self, h): return self.lm.lm_head(h) def _forward(self, ids: mx.array, cache: SeqCache, attn_mask, n_clean, last_only: bool, - rows: Optional[int] = None): - h = self.backbone.embed_tokens(ids) + rows: Optional[int] = None, embeds: Optional[mx.array] = None, mrope=None): + h = self.backbone.embed_tokens(ids) if embeds is None else embeds + offset = cache.pos recs = [] if n_clean is not None else None for i, layer in enumerate(self.backbone.layers): x = layer.input_layernorm(h) @@ -320,7 +348,7 @@ def _forward(self, ids: mx.array, cache: SeqCache, attn_mask, n_clean, last_only if recs is not None: recs.append(rec) else: - r = self._attention(layer.self_attn, x, c, attn_mask) + r = self._attention(layer.self_attn, x, c, attn_mask, offset, mrope) if recs is not None: recs.append(None) h = h + r @@ -333,26 +361,83 @@ def _forward(self, ids: mx.array, cache: SeqCache, attn_mask, n_clean, last_only return self._lm_head(h), recs # -- public forwards -------------------------------------------------------- + def _image_layout(self, cache: SeqCache, tokens: list): + """For a token list with image keys (negative ids): (3 x L positions, feature rows, + next position). Each image's key run must lie entirely inside ``tokens``.""" + m = self.vision.cfg.spatial_merge_size + pos3 = np.zeros((3, len(tokens)), dtype=np.int64) + feats, p, i = [], cache.pos, 0 + while i < len(tokens): + t = tokens[i] + if t >= 0: + pos3[:, i] = p + p += 1 + i += 1 + continue + ent = self.images.get(t) + if ent is None: + raise KeyError(f"image {t} is not registered") + gt, gh, gw = ent["grid"] + lt, lh, lw = gt, gh // m, gw // m + n = lt * lh * lw + if tokens[i : i + n] != [t] * n: + raise ValueError("an image's tokens must be prefilled in one call") + ti, hi, wi = np.meshgrid(np.arange(lt), np.arange(lh), np.arange(lw), indexing="ij") + pos3[:, i : i + n] = np.stack([ti.reshape(-1), hi.reshape(-1), wi.reshape(-1)]) + p + p += max(lt, lh, lw) + if ent.get("features") is None: + ent["features"] = self.vision(mx.array(ent["pixels"]), [ent["grid"]]) + mx.eval(ent["features"]) + feats.append(ent["features"]) + i += n + return pos3, (mx.concatenate(feats, axis=0) if feats else None), p + def prefill(self, cache: SeqCache, tokens: list) -> mx.array: """Causal prefill of ``tokens`` after whatever ``cache`` already holds. + Negative ids are image keys (see ``self.images``): their rows take the vision + features and the whole call uses multimodal RoPE positions. Returns the float32 logits [V] of the last token.""" assert cache.pending is None and tokens + has_image = any(t < 0 for t in tokens) + if has_image: + if self.vision is None: + raise ValueError("this checkpoint has no vision encoder") + pos3, feats, next_pos = self._image_layout(cache, tokens) + dims = int(self.args.head_dim * self.args.partial_rotary_factor) + section = (self.args.rope_scaling or {}).get("mrope_section", [11, 11, 10]) + feat_row = 0 logits = None for s in range(0, len(tokens), self.prefill_chunk): chunk = tokens[s : s + self.prefill_chunk] - ids = mx.array([chunk], dtype=mx.int32) mask = "causal" if len(chunk) > 1 else None - logits, _ = self._forward(ids, cache, mask, None, last_only=True) + if has_image: + real = [self.image_token_id if t < 0 else t for t in chunk] + ids = mx.array([real], dtype=mx.int32) + emb = self.backbone.embed_tokens(ids) + img_idx = [j for j, t in enumerate(chunk) if t < 0] + if img_idx: + rows = feats[feat_row : feat_row + len(img_idx)].astype(emb.dtype) + emb[0, mx.array(img_idx)] = rows + feat_row += len(img_idx) + cos, sin = mrope_cos_sin(pos3[:, s : s + len(chunk)], dims, self.args.rope_theta, section) + logits, _ = self._forward(ids, cache, mask, None, last_only=True, embeds=emb, mrope=(cos, sin)) + else: + ids = mx.array([chunk], dtype=mx.int32) + logits, _ = self._forward(ids, cache, mask, None, last_only=True) + cache.pos += len(chunk) cache.tokens.extend(chunk) if s + self.prefill_chunk < len(tokens): mx.eval([l.ssm if isinstance(l, GDNState) else l.keys for l in cache.layers]) if self.progress_cb is not None: self.progress_cb(len(chunk)) + if has_image: + cache.pos = next_pos return logits[0, -1].astype(mx.float32) def ar_step(self, cache: SeqCache, token: int) -> mx.array: logits, _ = self._forward(mx.array([[token]], dtype=mx.int32), cache, None, None, last_only=True) cache.tokens.append(token) + cache.pos += 1 return logits[0, -1].astype(mx.float32) def _canvas_mask(self, prefix_len: int, blk: int, n_clean: int) -> mx.array: @@ -399,5 +484,6 @@ def commit(self, cache: SeqCache, adv: int) -> None: rec.g[:, :adv], rec.beta[:, :adv], rec.prev.ssm) cache.layers[i] = GDNState(conv=conv, ssm=ssm) cache.tokens.extend(self._canvas_tokens[:adv]) + cache.pos += adv cache.pending = None cache.canvas_len = 0 diff --git a/inference/mlx/trida_mlx/profile.py b/inference/mlx/trida_mlx/profile.py index ec4518f..712f367 100644 --- a/inference/mlx/trida_mlx/profile.py +++ b/inference/mlx/trida_mlx/profile.py @@ -35,11 +35,11 @@ def timeit(fn, iters, warmup=3): def main(argv=None): ap = argparse.ArgumentParser() ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") - ap.add_argument("--gen-block", type=int, default=4) + ap.add_argument("--gen-block", type=int, default=None, help="N (canvas 2N-1); default 4") ap.add_argument("--iters", type=int, default=20) a = ap.parse_args(argv) eng = Engine(a.model, gen_block=a.gen_block, prompt_cache=False) - rt, n, mask = eng.rt, a.gen_block, eng.mask_id + rt, n, mask = eng.rt, eng.n, eng.mask_id blk = 2 * n - 1 prompt = eng.encode(eng.render(LONG_PROMPT, enable_thinking=False)) cache = rt.make_cache() @@ -83,7 +83,7 @@ def run_ar(): def run_canvas(rows_=None, adv=n): restore() - L = rt.canvas(cache, [11, 12, 13, 14][:n] + [mask] * (n - 1), n, rows=rows_) + L = rt.canvas(cache, list(range(11, 11 + n)) + [mask] * (n - 1), n, rows=rows_) am = mx.argmax(L, axis=-1) mx.eval(am) rt.commit(cache, adv) @@ -102,7 +102,7 @@ def run_canvas(rows_=None, adv=n): def build_only(): restore() t = time.perf_counter() - L = rt.canvas(cache, [11, 12, 13, 14][:n] + [mask] * (n - 1), n) + L = rt.canvas(cache, list(range(11, 11 + n)) + [mask] * (n - 1), n) dt = time.perf_counter() - t mx.eval(L) cache.pending = None diff --git a/inference/mlx/trida_mlx/server.py b/inference/mlx/trida_mlx/server.py index 7c26356..0928cef 100644 --- a/inference/mlx/trida_mlx/server.py +++ b/inference/mlx/trida_mlx/server.py @@ -18,6 +18,7 @@ import queue import sys import threading +from typing import Optional import time import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -40,27 +41,64 @@ def _sampling(body: dict) -> SamplingParams: ) -def _text_of(content) -> str: - if isinstance(content, list): - return "".join(p.get("text", "") for p in content if isinstance(p, dict) and "text" in p) - return "" if content is None else str(content) - - -def _normalize_messages(messages: list) -> list: - """OpenAI -> the Trida (Qwen3.5) chat template's contract: - * content parts -> text (no vision here), None -> "" +def _image_source(part: dict): + """The image reference of an OpenAI / Responses / plain content part, or None.""" + if not isinstance(part, dict): + return None + t = part.get("type") + if t in ("image_url", "input_image") or "image_url" in part: + u = part.get("image_url") + return u.get("url") if isinstance(u, dict) else u + if t == "image" or "image" in part: + u = part.get("image") + return u.get("url") if isinstance(u, dict) else u + return None + + +def _content(content, images: Optional[list]): + """Template content for one message. With ``images`` (a list to append sources to) image + parts stay as ``{"type": "image"}`` items so the chat template emits its vision + placeholder; without it they become a short text marker (no vision encoder).""" + if content is None: + return "" + if not isinstance(content, list): + return str(content) + items, has_image = [], False + for p in content: + src = _image_source(p) + if src is not None: + if images is not None: + images.append(src) + items.append({"type": "image"}) + has_image = True + else: + items.append({"type": "text", "text": "[image omitted: this model has no vision encoder]"}) + elif isinstance(p, dict) and "text" in p: + items.append({"type": "text", "text": p.get("text") or ""}) + if not has_image: + return "".join(i["text"] for i in items) + return items + + +def _normalize_messages(messages: list, vision: bool = False): + """OpenAI -> the Trida (Qwen3.5) chat template's contract. Returns (messages, image_sources). + * text parts joined; image parts kept as template image items (in order) when ``vision`` * ``developer`` -> ``system``; all leading system messages merged into one; a system message later in the conversation (harness nudges, compaction notes) becomes a user turn, because - the template only accepts a system message first + the template only accepts a system message first (images in system messages are dropped) * assistant ``tool_calls[].function.arguments`` as dicts (the template iterates them) """ - out, sys_parts, seen_non_system = [], [], False + out, sys_parts, seen_non_system, images = [], [], False, [] for m in messages: m = dict(m) role = m.get("role") if role == "developer": role = m["role"] = "system" - m["content"] = _text_of(m.get("content")) + if role in ("system", "assistant"): + c = _content(m.get("content"), None) + m["content"] = c if isinstance(c, str) else "" + else: + m["content"] = _content(m.get("content"), images if vision else None) if role == "system": if not seen_non_system: sys_parts.append(m["content"]) @@ -86,7 +124,7 @@ def _normalize_messages(messages: list) -> list: out.append(m) if sys_parts: out.insert(0, {"role": "system", "content": "\n\n".join(p for p in sys_parts if p)}) - return out + return out, images def _thinking(body: dict) -> bool: @@ -107,6 +145,10 @@ class ContextOverflow(Exception): pass +class BadRequest(Exception): + pass + + class _StopMatcher: def __init__(self, stops): self.stops = [s for s in ([stops] if isinstance(stops, str) else (stops or [])) if s] @@ -245,7 +287,9 @@ def do_OPTIONS(self): def _model_card(self): ctx = DEFAULTS["context_length"] return {"id": DEFAULTS["served_name"], "object": "model", "owned_by": "trillionlabs", "created": 0, - "context_length": ctx, "max_model_len": ctx, "root": str(ENGINE.path)} + "context_length": ctx, "max_model_len": ctx, "root": str(ENGINE.path), + "supports_vision": ENGINE.supports_vision, + "input_modalities": ["text", "image"] if ENGINE.supports_vision else ["text"]} def do_GET(self): path = self.path.split("?")[0].rstrip("/") @@ -277,6 +321,8 @@ def do_POST(self): if route in ("/v1/completions", "/completions"): return self._completion(body) return self._json(404, {"error": {"message": f"no route {self.path}"}}) + except BadRequest as e: + return self._json(400, {"error": {"message": str(e), "type": "invalid_request_error"}}) except ContextOverflow as e: return self._json(400, {"error": {"message": str(e), "type": "invalid_request_error", "param": "messages", "code": "context_length_exceeded"}}) @@ -331,11 +377,15 @@ def _usage(self, stats: DecodeStats): "trida": stats.as_dict()} def _chat(self, body: dict): - messages = _normalize_messages(body.get("messages") or []) + messages, image_srcs = _normalize_messages(body.get("messages") or [], vision=ENGINE.supports_vision) tools = body.get("tools") if body.get("tool_choice") != "none" else None kw = {k: v for k, v in (body.get("chat_template_kwargs") or {}).items() if k != "enable_thinking"} prompt_text = ENGINE.render(messages, tools=tools, enable_thinking=_thinking(body), **kw) - prompt = ENGINE.encode(prompt_text) + try: + images = [ENGINE.add_image(src) for src in image_srcs] + except Exception as e: # noqa: BLE001 + raise BadRequest(f"could not load image: {type(e).__name__}: {e}") + prompt = ENGINE.encode(prompt_text, images) self._budget(prompt, body) # 400 before any streaming starts started_in_think = prompt_text.rstrip().endswith("") rid, created = f"chatcmpl-{uuid.uuid4().hex[:24]}", int(time.time()) @@ -421,7 +471,7 @@ def main(argv=None): ap.add_argument("--host", default="127.0.0.1") ap.add_argument("--port", type=int, default=8080) ap.add_argument("--mode", default="self-spec", choices=["self-spec", "causal"]) - ap.add_argument("--gen-block", type=int, default=4, help="N: tokens per self-spec step (canvas 2N-1)") + ap.add_argument("--gen-block", type=int, default=None, help="N tokens per self-spec step (canvas 2N-1); default 4 (fastest on Apple Silicon)") ap.add_argument("--prefill-chunk", type=int, default=512) ap.add_argument("--no-prompt-cache", action="store_true") ap.add_argument("--temperature", type=float, default=1.0) @@ -432,6 +482,8 @@ def main(argv=None): help="advertised + enforced context window (Hermes needs >= 64000)") ap.add_argument("--served-model-name", default="trida2.0-4b") ap.add_argument("--log-requests", default=None, metavar="DIR", help="dump every POST body as JSON (debugging)") + ap.add_argument("--max-image-pixels", type=int, default=1024 * 1024, + help="resize budget per image (every 32x32 px = 1 LM token; 1024*1024 -> <=1024 tokens)") ap.add_argument("--cache-slots", type=int, default=3, help="conversations whose prompt cache is kept (main loop + side requests)") ap.add_argument("--no-think", action="store_true", help="default enable_thinking=false") @@ -443,13 +495,14 @@ def main(argv=None): print(f"[trida-mlx] loading {a.model} ...", flush=True) WORKER = _Worker(dict(model=a.model, mode=a.mode, gen_block=a.gen_block, prefill_chunk=a.prefill_chunk, prompt_cache=not a.no_prompt_cache, fused_gdn=not a.no_fused_gdn, - cache_slots=a.cache_slots)) + cache_slots=a.cache_slots, max_image_pixels=a.max_image_pixels)) WORKER.start() WORKER.ready.wait() if WORKER.error is not None: raise WORKER.error - print(f"[trida-mlx] ready in {ENGINE.load_s:.1f}s mode={a.mode} N={a.gen_block} mask_id={ENGINE.mask_id} " - f"eos={sorted(ENGINE.eos_ids)} ctx={a.context_length} model={a.served_model_name!r} -> " + print(f"[trida-mlx] ready in {ENGINE.load_s:.1f}s mode={a.mode} N={ENGINE.n} mask_id={ENGINE.mask_id} " + f"eos={sorted(ENGINE.eos_ids)} ctx={a.context_length} vision={ENGINE.supports_vision} " + f"model={a.served_model_name!r} -> " f"http://{a.host}:{a.port}/v1", flush=True) srv = ThreadingHTTPServer((a.host, a.port), Handler) try: diff --git a/inference/mlx/trida_mlx/verify.py b/inference/mlx/trida_mlx/verify.py index bb45629..49a57bf 100644 --- a/inference/mlx/trida_mlx/verify.py +++ b/inference/mlx/trida_mlx/verify.py @@ -27,19 +27,32 @@ def main(argv=None): ap = argparse.ArgumentParser() ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") ap.add_argument("--max-tokens", type=int, default=256) - ap.add_argument("--gen-block", type=int, default=4) + ap.add_argument("--gen-block", type=int, default=None, help="N (canvas 2N-1); default 4") + ap.add_argument("--image", default=None, help="also run the checks on an image prompt (path or URL)") a = ap.parse_args(argv) eng = Engine(a.model, gen_block=a.gen_block, prompt_cache=False) - rt, n = eng.rt, a.gen_block + rt, n = eng.rt, eng.n args = rt.args n_lin = sum(l.is_linear for l in rt.backbone.layers) print(f"model {a.model}: {args.num_hidden_layers} layers ({n_lin} gated-delta, {args.num_hidden_layers - n_lin} " f"full-attn), hidden {args.hidden_size}, vocab {args.vocab_size}, tied={args.tie_word_embeddings}, " f"quant={eng.cfg.get('quantization')}") print(f"mask_id={eng.mask_id} ({eng.tokenizer.convert_ids_to_tokens(eng.mask_id)!r}) eos={sorted(eng.eos_ids)} " - f"load {eng.load_s:.1f}s device={mx.default_device()}") + f"load {eng.load_s:.1f}s device={mx.default_device()} N={n} (canvas {2 * n - 1}) " + f"vision={'yes' if eng.supports_vision else 'no'}") - prompt = eng.encode(eng.render(PROMPT, enable_thinking=False)) + if a.image: + if not eng.supports_vision: + raise SystemExit("--image given but this checkpoint has no vision encoder") + t = time.perf_counter() + key, n_img = eng.add_image(a.image) + msgs = [{"role": "user", "content": [{"type": "image"}, + {"type": "text", "text": "Describe this image in two sentences."}]}] + prompt = eng.encode(eng.render(msgs, enable_thinking=False), [(key, n_img)]) + print(f"[image] {a.image}: grid {eng.rt.images[key]['grid']} -> {n_img} tokens " + f"(preprocess {time.perf_counter() - t:.2f}s)") + else: + prompt = eng.encode(eng.render(PROMPT, enable_thinking=False)) sp = SamplingParams(temperature=0.0) # 1) AR reference with margins diff --git a/inference/mlx/trida_mlx/vision.py b/inference/mlx/trida_mlx/vision.py new file mode 100644 index 0000000..3e2e91b --- /dev/null +++ b/inference/mlx/trida_mlx/vision.py @@ -0,0 +1,336 @@ +"""Vision input for Trida checkpoints that ship a Qwen3.5 (Qwen3-VL family) vision encoder. + +Three pieces, all MLX / numpy / PIL only: + +* ``preprocess`` Qwen2/3-VL image preprocessing (smart resize to a multiple of + patch*merge, normalize, flatten into (C, T, P, P) patches) -> pixel rows + + ``grid_thw``. +* ``VisionTower`` the Qwen3-VL vision encoder: Conv3d patch embed, bilinearly + interpolated learned position embedding, 2-D rotary ViT blocks and the + 2x2 patch merger/projector into the LM hidden size. + Adapted from mlx-vlm's ``qwen3_vl/vision.py`` (MIT, Apple Inc. and + contributors) so the parameter names match HF/mlx-vlm checkpoints. +* ``mrope_*`` interleaved multimodal RoPE positions for prompts that contain images + (text tokens get equal t/h/w positions, i.e. ordinary RoPE). +""" +from __future__ import annotations + +import base64 +import hashlib +import io +import math +from dataclasses import dataclass, field +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + + +# ----------------------------------------------------------------------------- config +@dataclass +class VisionConfig: + depth: int = 24 + hidden_size: int = 1024 + intermediate_size: int = 4096 + num_heads: int = 16 + patch_size: int = 16 + temporal_patch_size: int = 2 + spatial_merge_size: int = 2 + in_channels: int = 3 + out_hidden_size: int = 2560 + num_position_embeddings: int = 2304 + deepstack_visual_indexes: list = field(default_factory=list) + + @classmethod + def from_dict(cls, d: dict) -> "VisionConfig": + keys = cls.__dataclass_fields__.keys() + return cls(**{k: v for k, v in d.items() if k in keys}) + + +@dataclass +class ImageProcessorConfig: + patch_size: int = 16 + temporal_patch_size: int = 2 + merge_size: int = 2 + min_pixels: int = 64 * 32 * 32 + max_pixels: int = 1024 * 1024 + image_mean: tuple = (0.5, 0.5, 0.5) + image_std: tuple = (0.5, 0.5, 0.5) + + @classmethod + def from_files(cls, path, max_pixels: Optional[int] = None) -> "ImageProcessorConfig": + import json + c = cls() + for name in ("preprocessor_config.json", "processor_config.json"): + f = path / name + if not f.is_file(): + continue + d = json.loads(f.read_text()) + d = d.get("image_processor", d) + c.patch_size = int(d.get("patch_size", c.patch_size)) + c.temporal_patch_size = int(d.get("temporal_patch_size", c.temporal_patch_size)) + c.merge_size = int(d.get("merge_size", c.merge_size)) + size = d.get("size") or {} + c.min_pixels = int(size.get("shortest_edge", d.get("min_pixels", c.min_pixels))) + c.max_pixels = int(size.get("longest_edge", d.get("max_pixels", c.max_pixels))) + c.image_mean = tuple(d.get("image_mean", c.image_mean)) + c.image_std = tuple(d.get("image_std", c.image_std)) + break + if max_pixels: # on-device budget: every 32x32 pixels is one LM token + c.max_pixels = min(c.max_pixels, max_pixels) + c.min_pixels = min(c.min_pixels, c.max_pixels) + return c + + +# ----------------------------------------------------------------------------- preprocessing +def load_image(src): + """PIL image from bytes, a data: URL, an http(s) URL or a local path.""" + from PIL import Image + + if isinstance(src, (bytes, bytearray)): + data = bytes(src) + elif isinstance(src, str) and src.startswith("data:"): + data = base64.b64decode(src.split(",", 1)[1]) + elif isinstance(src, str) and src.startswith(("http://", "https://")): + import urllib.request + + with urllib.request.urlopen(src, timeout=20) as r: + data = r.read() + elif isinstance(src, str): + with open(src.removeprefix("file://"), "rb") as fh: + data = fh.read() + else: + raise ValueError(f"unsupported image source {type(src)}") + img = Image.open(io.BytesIO(data)) + img.load() + return img.convert("RGB"), data + + +def smart_resize(height: int, width: int, factor: int, min_pixels: int, max_pixels: int): + if max(height, width) / max(1, min(height, width)) > 200: + raise ValueError("image aspect ratio must be < 200") + h_bar = max(factor, round(height / factor) * factor) + w_bar = max(factor, round(width / factor) * factor) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = max(factor, math.floor(height / beta / factor) * factor) + w_bar = max(factor, math.floor(width / beta / factor) * factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = math.ceil(height * beta / factor) * factor + w_bar = math.ceil(width * beta / factor) * factor + return h_bar, w_bar + + +def preprocess(img, cfg: ImageProcessorConfig): + """-> (pixel_rows float32 [N, C*T*P*P], grid_thw (t, h, w)).""" + from PIL import Image + + P, T, m = cfg.patch_size, cfg.temporal_patch_size, cfg.merge_size + h, w = smart_resize(img.height, img.width, P * m, cfg.min_pixels, cfg.max_pixels) + x = np.asarray(img.resize((w, h), Image.BICUBIC), dtype=np.float32) / 255.0 # H W C + x = (x - np.array(cfg.image_mean, np.float32)) / np.array(cfg.image_std, np.float32) + x = x.transpose(2, 0, 1)[None] # 1 C H W + x = np.repeat(x, T, axis=0) # a still image fills the temporal patch + C = x.shape[1] + gt, gh, gw = 1, h // P, w // P + x = x.reshape(gt, T, C, gh // m, m, P, gw // m, m, P) + x = x.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8) + return x.reshape(gt * gh * gw, C * T * P * P), (gt, gh, gw) + + +def image_key(data: bytes, cfg: ImageProcessorConfig) -> int: + """Stable negative id standing for one image's tokens in cache keys (never a real token).""" + h = hashlib.sha1(data + f"{cfg.max_pixels}/{cfg.min_pixels}".encode()).hexdigest() + return -(int(h[:12], 16) % (1 << 46)) - 1 + + +# ----------------------------------------------------------------------------- encoder +def _rotate_half(x): + half = x.shape[-1] // 2 + return mx.concatenate([-x[..., half:], x[..., :half]], axis=-1) + + +class PatchEmbed(nn.Module): + def __init__(self, cfg: VisionConfig): + super().__init__() + self.cfg = cfg + k = [cfg.temporal_patch_size, cfg.patch_size, cfg.patch_size] + self.proj = nn.Conv3d(cfg.in_channels, cfg.hidden_size, kernel_size=k, stride=k, bias=True) + + def __call__(self, x): + c = self.cfg + x = x.reshape(-1, c.in_channels, c.temporal_patch_size, c.patch_size, c.patch_size).moveaxis(1, 4) + return self.proj(x).reshape(-1, c.hidden_size) + + +class _Attention(nn.Module): + def __init__(self, dim: int, num_heads: int): + super().__init__() + self.num_heads = num_heads + self.scale = (dim // num_heads) ** -0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim) + + def __call__(self, x, cu_seqlens: list, cos, sin): + L = x.shape[0] + q, k, v = mx.split(self.qkv(x).reshape(L, 3, self.num_heads, -1).transpose(1, 0, 2, 3), 3) + q, k, v = q[0], k[0], v[0] # [L, H, D] + dt = q.dtype + q = (q.astype(mx.float32) * cos + _rotate_half(q.astype(mx.float32)) * sin).astype(dt) + k = (k.astype(mx.float32) * cos + _rotate_half(k.astype(mx.float32)) * sin).astype(dt) + q, k, v = (t.transpose(1, 0, 2)[None] for t in (q, k, v)) # [1, H, L, D] + outs = [] + for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]): # attention within each image + outs.append(mx.fast.scaled_dot_product_attention( + q[:, :, s:e], k[:, :, s:e], v[:, :, s:e], scale=self.scale)) + o = mx.concatenate(outs, axis=2) + return self.proj(o[0].transpose(1, 0, 2).reshape(L, -1)) + + +class _MLP(nn.Module): + def __init__(self, dim, hidden): + super().__init__() + self.linear_fc1 = nn.Linear(dim, hidden, bias=True) + self.linear_fc2 = nn.Linear(hidden, dim, bias=True) + + def __call__(self, x): + return self.linear_fc2(nn.gelu_approx(self.linear_fc1(x))) + + +class _Block(nn.Module): + def __init__(self, cfg: VisionConfig): + super().__init__() + self.norm1 = nn.LayerNorm(cfg.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(cfg.hidden_size, eps=1e-6) + self.attn = _Attention(cfg.hidden_size, cfg.num_heads) + self.mlp = _MLP(cfg.hidden_size, cfg.intermediate_size) + + def __call__(self, x, cu, cos, sin): + x = x + self.attn(self.norm1(x), cu, cos, sin) + return x + self.mlp(self.norm2(x)) + + +class _Merger(nn.Module): + def __init__(self, cfg: VisionConfig): + super().__init__() + self.hidden = cfg.hidden_size * cfg.spatial_merge_size ** 2 + self.norm = nn.LayerNorm(cfg.hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden, self.hidden) + self.linear_fc2 = nn.Linear(self.hidden, cfg.out_hidden_size) + + def __call__(self, x): + x = self.norm(x).reshape(-1, self.hidden) + return self.linear_fc2(nn.gelu(self.linear_fc1(x))) + + +class VisionTower(nn.Module): + def __init__(self, cfg: VisionConfig): + super().__init__() + if cfg.deepstack_visual_indexes: + raise NotImplementedError("deepstack vision features are not supported") + self.cfg = cfg + self.patch_embed = PatchEmbed(cfg) + self.pos_embed = nn.Embedding(cfg.num_position_embeddings, cfg.hidden_size) + self.num_grid_per_side = int(cfg.num_position_embeddings ** 0.5) + self.blocks = [_Block(cfg) for _ in range(cfg.depth)] + self.merger = _Merger(cfg) + + # 2-D rotary (rows, cols of the full-resolution patch grid, in merge-window order) + def _rotary(self, grid): + m = self.cfg.spatial_merge_size + head_dim = self.cfg.hidden_size // self.cfg.num_heads + dim = head_dim // 2 + inv = 1.0 / (10000.0 ** (np.arange(0, dim, 2, dtype=np.float32) / dim)) + rows, cols = [], [] + for t, h, w in grid: + r = (np.arange(h // m)[:, None, None, None] * m + np.arange(m)[None, None, :, None]) + c = (np.arange(w // m)[None, :, None, None] * m + np.arange(m)[None, None, None, :]) + r = np.broadcast_to(r, (h // m, w // m, m, m)).reshape(-1) + c = np.broadcast_to(c, (h // m, w // m, m, m)).reshape(-1) + rows.append(np.tile(r, t)); cols.append(np.tile(c, t)) + r, c = np.concatenate(rows), np.concatenate(cols) + freqs = np.concatenate([np.outer(r, inv), np.outer(c, inv)], axis=-1) # [L, head_dim/2] + emb = np.concatenate([freqs, freqs], axis=-1)[:, None, :] # [L, 1, head_dim] + return mx.array(np.cos(emb)), mx.array(np.sin(emb)) + + def _pos_embed(self, grid): + n, m = self.num_grid_per_side, self.cfg.spatial_merge_size + outs = [] + for t, h, w in grid: + hi, wi = np.linspace(0, n - 1, h), np.linspace(0, n - 1, w) + hf, wf = hi.astype(np.int32), wi.astype(np.int32) + hc, wc = np.minimum(hf + 1, n - 1), np.minimum(wf + 1, n - 1) + dh, dw = (hi - hf)[:, None], (wi - wf)[None, :] + idx = [hf[:, None] * n + wf[None], hf[:, None] * n + wc[None], + hc[:, None] * n + wf[None], hc[:, None] * n + wc[None]] + wts = [(1 - dh) * (1 - dw), (1 - dh) * dw, dh * (1 - dw), dh * dw] + e = None + for i, wt in zip(idx, wts): + term = self.pos_embed(mx.array(i.reshape(-1))) * mx.array(wt.reshape(-1, 1).astype(np.float32)) + e = term if e is None else e + term + D = e.shape[-1] + e = mx.tile(e, (t, 1)).reshape(t, h // m, m, w // m, m, D).transpose(0, 1, 3, 2, 4, 5).reshape(-1, D) + outs.append(e) + return mx.concatenate(outs, axis=0) + + def __call__(self, pixels: mx.array, grid: list) -> mx.array: + """pixels [N, C*T*P*P], grid [(t, h, w), ...] -> merged features [N / merge^2, out_hidden].""" + x = self.patch_embed(pixels.astype(self.patch_embed.proj.weight.dtype)) + x = x + self._pos_embed(grid).astype(x.dtype) + cos, sin = self._rotary(grid) + cu = [0] + for t, h, w in grid: + for _ in range(t): + cu.append(cu[-1] + h * w) + for blk in self.blocks: + x = blk(x, cu, cos, sin) + return self.merger(x) + + +def sanitize_vision_weights(weights: dict) -> dict: + """{any HF/mlx-vlm vision key: array} -> {VisionTower key: array}.""" + out = {} + for k, v in weights.items(): + for pre in ("model.visual.", "visual.", "vision_tower.", "model.language_model.visual."): + if k.startswith(pre): + k = k[len(pre):] + break + else: + continue + if "position_ids" in k: + continue + if k == "patch_embed.proj.weight" and v.ndim == 5 and v.shape[1] in (1, 3) and v.shape[-1] not in (1, 3): + v = v.transpose(0, 2, 3, 4, 1) # PyTorch NCDHW -> MLX NDHWC + out[k] = v + return out + + +def is_vision_key(k: str) -> bool: + return k.startswith(("model.visual.", "visual.", "vision_tower.", "model.language_model.visual.")) + + +# ----------------------------------------------------------------------------- mrope +def mrope_cos_sin(pos3: np.ndarray, dims: int, base: float, section) -> tuple: + """pos3 [3, L] (t, h, w) -> cos, sin [L, dims] for interleaved MRoPE (Qwen3-VL/3.5): + frequency j takes the h position if j % 3 == 1 and j < 3*section[1], the w position if + j % 3 == 2 and j < 3*section[2], else the t position.""" + inv = 1.0 / (base ** (np.arange(0, dims, 2, dtype=np.float64) / dims)) # [dims/2] + n = inv.shape[0] + axis = np.zeros(n, dtype=np.int64) + axis[1: section[1] * 3: 3] = 1 + axis[2: section[2] * 3: 3] = 2 + pos = pos3.astype(np.float64)[axis] # [n, L] + freqs = (pos * inv[:, None]).T # [L, n] + emb = np.concatenate([freqs, freqs], axis=-1) + return mx.array(np.cos(emb).astype(np.float32)), mx.array(np.sin(emb).astype(np.float32)) + + +def apply_mrope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array: + """x [B, H, L, D]; rotates the first cos.shape[-1] dims (non-traditional / rotate-half).""" + d = cos.shape[-1] + xr, xp = x[..., :d], x[..., d:] + xf = xr.astype(mx.float32) + out = (xf * cos + _rotate_half(xf) * sin).astype(x.dtype) + return mx.concatenate([out, xp], axis=-1) if xp.shape[-1] else out diff --git a/inference/mlx/uv.lock b/inference/mlx/uv.lock index 89282bf..b8f5d4d 100644 --- a/inference/mlx/uv.lock +++ b/inference/mlx/uv.lock @@ -436,6 +436,40 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, ] +[[package]] +name = "pillow" +version = "12.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6b/ba/3762f376a2948e3036488d773a146e0ae6ecc2ca03ac20e2615bd0b2ba02/pillow-12.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:bcc33feacfaefce60c12fd500a277533bdc02b10a19f7f6d348763d8140bbba7", size = 4785287, upload-time = "2026-07-01T11:53:29.761Z" }, + { url = "https://files.pythonhosted.org/packages/07/50/b5d688cc9c52d4482f3d5bcab6ce20bc2a74a85d2343841c907444a3be2c/pillow-12.3.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5594fc43d548a7ed94949d139aa1341b270f1863f11cfd37f5a6c8b778a6b67f", size = 6253754, upload-time = "2026-07-01T11:53:32.298Z" }, + { url = "https://files.pythonhosted.org/packages/4e/89/36f4cd76cf4baf05c50ababb976249153f18c959171c7f6ba09a6f217260/pillow-12.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f0606c8bf2cdefea14a43530f7657cbbb7ecf1c4222512492ef4a4434a9501ec", size = 6925605, upload-time = "2026-07-01T11:53:34.487Z" }, + { url = "https://files.pythonhosted.org/packages/eb/c0/4de58cf6633b9e3a6061ef4be6fb91fc3c90b812ece886f531e3c523d777/pillow-12.3.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:85f998ea1848bc6757289e739cfbdda3a04adfd58b02fc018ce54d754a5ce468", size = 6327788, upload-time = "2026-07-01T11:53:36.433Z" }, + { url = "https://files.pythonhosted.org/packages/87/3c/14d53682a19550dbbaf3b598f807d5457646c510805a44c7d7891cd1cd1a/pillow-12.3.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:25b9b82bb22e6e2b3cd07b39c68b7b862001226cb3dff7130d1cb914121b39ed", size = 7036288, upload-time = "2026-07-01T11:53:38.712Z" }, + { url = "https://files.pythonhosted.org/packages/b2/5b/a02d30018abd97ced9f5a6c63d28597694a00d066516b9c1c6de45859fc9/pillow-12.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:37d6d0a00072fd2948eb22bce7e1475f34569d90c87c59f7a2ec59541b77f7a6", size = 4785266, upload-time = "2026-07-01T11:53:49.079Z" }, + { url = "https://files.pythonhosted.org/packages/c8/98/766667a4be768150a202836acd9fad19c06824ca86c4286d3cf6b274964e/pillow-12.3.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bcb46e2f9feff8d06323983bd83ed00c201fdcab3d74973e7072a889b3979fcd", size = 6263814, upload-time = "2026-07-01T11:53:51.32Z" }, + { url = "https://files.pythonhosted.org/packages/3b/2d/ede717bc1144f63886c21fd349bb95860b0d1a21149ff16f2bb362b612b6/pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23d27a3e0307ec2244cc51e7287b919aa68d097504ebe19df4e76a98a3eea5bd", size = 6934408, upload-time = "2026-07-01T11:53:53.487Z" }, + { url = "https://files.pythonhosted.org/packages/a3/48/9c58b685e69d49c31af6c8eb9012055fab7e665785165c84796e2c73ce72/pillow-12.3.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4f883547d4b7f0495ebe7056b0cc2aea76094e7a4abc8e933540f3271df27d9c", size = 6337160, upload-time = "2026-07-01T11:53:55.457Z" }, + { url = "https://files.pythonhosted.org/packages/ff/fa/dc2a5c0ba6df93f67c31d34b808b7ce440b40cdbf96f0b81cde1d1e6fa93/pillow-12.3.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:236ff70b9312fb68943c703aa842ca6a758abfa45ac187a5e7c1452e96ef72b5", size = 7045172, upload-time = "2026-07-01T11:53:57.736Z" }, + { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" }, + { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" }, + { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" }, + { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, + { url = "https://files.pythonhosted.org/packages/9d/ac/31fb64e1e7efb5a4b50cd3d92049ba89ac6e4d8d3bb6a74e15048ca3353e/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:21900ce7ba264168cd50defae43cd75d25c833ad4ad6e73ffc5596d12e25ac89", size = 4161684, upload-time = "2026-07-01T11:54:25.934Z" }, + { url = "https://files.pythonhosted.org/packages/87/b4/9805e23d2b4d77842b468513841fda254ee42f0289d25088340e4ff46e2d/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:4e8c2a84d977f50b9daed6eeaf3baef67d00d5d74d932288f02cb94518ee3ace", size = 4255487, upload-time = "2026-07-01T11:54:27.935Z" }, + { url = "https://files.pythonhosted.org/packages/df/39/ecf519435a200c693fe053a6ee4d835b41cf963a4dfc2551c4e637cb2a71/pillow-12.3.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:ae26d61dfa7a47befdc7572b521024e8745f3d809bd95ca9505a7bba9ef849ec", size = 3696433, upload-time = "2026-07-01T11:54:29.813Z" }, + { url = "https://files.pythonhosted.org/packages/10/76/8803c13605b763d33d156c4678fc77f8443389c0c51c8aef707bb02015f4/pillow-12.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d69141514cc30b774ceea5e3ed3a6635c8d8a96edf664689b890f4089111fb35", size = 4780109, upload-time = "2026-07-01T11:54:34.026Z" }, + { url = "https://files.pythonhosted.org/packages/1f/01/e18aff37cb0b4aac47ac90f016d347a49aca667ef97f190b06ac2aabc928/pillow-12.3.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f7401aebd7f581d7f83a439d87d474999317ee099218e5ad25d125290990ba65", size = 6263736, upload-time = "2026-07-01T11:54:36.131Z" }, + { url = "https://files.pythonhosted.org/packages/f7/62/de5bdd77d935331f4f802edc11e4d82950f642caad6cb2f949837b8560e2/pillow-12.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0847a763afefb695bc912d7c131e7e0632d4edc1d8698f58ddabec8e46b8b6d3", size = 6937129, upload-time = "2026-07-01T11:54:38.216Z" }, + { url = "https://files.pythonhosted.org/packages/70/4d/105627a13300c5e0df1d174230b32fd1273062c96f7745fd552b945d1e1d/pillow-12.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:571b9fcb07b97ef3a492028fb3d2dc0993ca23a06138b0315286566d29ef718a", size = 6339562, upload-time = "2026-07-01T11:54:40.354Z" }, + { url = "https://files.pythonhosted.org/packages/6b/1d/f13de01a553988ab895ba1c722e06cf3144d4f57656fd5b81b6d881f1179/pillow-12.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:756c768d0c9c2955feb7a56c37ea24aea2e369f8d36a88da270b6a9f19e62b5e", size = 7049439, upload-time = "2026-07-01T11:54:42.489Z" }, + { url = "https://files.pythonhosted.org/packages/46/3e/51fabf59d5ab801ceab709453d3ab6b180083496579549de4c45ced6528a/pillow-12.3.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:b343699e8308bdc51978310e1c959c584e7869cc8c40780058c87da7781a1e94", size = 4736058, upload-time = "2026-07-01T11:56:28.041Z" }, + { url = "https://files.pythonhosted.org/packages/bf/20/22fe9384b7949e25fb1293bcfc84fb82590ff4ea6b37c95b24d26d793d86/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fbd139c8447d25dd750ab79ee274cc5e1fe80fc56340ab10b18a195e1b6eca3e", size = 5237776, upload-time = "2026-07-01T11:56:30.263Z" }, + { url = "https://files.pythonhosted.org/packages/08/14/f6ba68107680ffa74b39985f3f30884e41318fbc4250caa423c79b4788bb/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7e480451b9fa137494bccd3a7d69adbe8ac65a87d97be61e11f1b1050a5bac3", size = 5860358, upload-time = "2026-07-01T11:56:32.68Z" }, +] + [[package]] name = "protobuf" version = "7.36.2" @@ -704,6 +738,7 @@ dependencies = [ { name = "mlx" }, { name = "mlx", extra = ["cpu"], marker = "sys_platform == 'linux'" }, { name = "mlx-lm" }, + { name = "pillow" }, { name = "transformers" }, ] @@ -713,6 +748,7 @@ requires-dist = [ { name = "mlx", marker = "sys_platform == 'darwin'", specifier = ">=0.29" }, { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = ">=0.29" }, { name = "mlx-lm", specifier = ">=0.31" }, + { name = "pillow" }, { name = "transformers", specifier = ">=4.57" }, ] From e103b8202b0e312ef91eff95b355f8b6a165c880 Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Tue, 6 Oct 2026 06:39:20 +0000 Subject: [PATCH 6/9] perf(inference/mlx): long-context profiling and real-prompt benchmarks - profile: --context N profiles at agent-length prefixes (14K-33K); adds timing-only bounds for a cheaper canvas attention mask (no mask, built-in causal) to size any mask-path optimization. - bench: --prompts takes a JSONL of OpenAI chat requests (e.g. logged Hermes requests), normalized like the server, images included. GQA head folding for canvas attention was tried and measured no gain on M3 Pro at 28K context (68.29 vs 67.83 ms), so it is not included. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- inference/mlx/trida_mlx/bench.py | 12 +++++++++--- inference/mlx/trida_mlx/profile.py | 19 +++++++++++++++++-- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/inference/mlx/trida_mlx/bench.py b/inference/mlx/trida_mlx/bench.py index b0518cb..2a23660 100644 --- a/inference/mlx/trida_mlx/bench.py +++ b/inference/mlx/trida_mlx/bench.py @@ -21,6 +21,7 @@ from .decode import DecodeStats, SamplingParams from .engine import Engine +from .server import _normalize_messages BUILTIN = [ {"name": "gsm8k", "messages": [{"role": "user", "content": @@ -94,9 +95,14 @@ def main(argv=None): pass results = [] for it in items: - text = it.get("prompt") or eng.render(it["messages"], tools=it.get("tools"), - enable_thinking=not a.no_think) - pids = eng.encode(text) + images = [] + if it.get("prompt"): + text = it["prompt"] + else: # OpenAI-style messages (e.g. captured with server --log-requests), normalized like the server + msgs, srcs = _normalize_messages(it["messages"], vision=eng.supports_vision) + images = [eng.add_image(s) for s in srcs] + text = eng.render(msgs, tools=it.get("tools"), enable_thinking=not a.no_think) + pids = eng.encode(text, images) row = {"name": it.get("name", f"p{len(results)}"), "prompt_tokens": len(pids), "runs": {}} outs = {} for m in modes: diff --git a/inference/mlx/trida_mlx/profile.py b/inference/mlx/trida_mlx/profile.py index 712f367..5f428de 100644 --- a/inference/mlx/trida_mlx/profile.py +++ b/inference/mlx/trida_mlx/profile.py @@ -5,7 +5,7 @@ Times, at the same KV length: the AR step (1 token) vs the self-spec canvas step (2N-1 tokens), each split into LM head / MLPs / full attention / gated-delta, plus the Python graph-build time and the commit, and A/Bs the step optimizations (fused canvas GDN -kernel, LM-head row skipping on cold starts). Medians over --iters runs, GPU-synchronized. +kernel, LM-head row skipping on cold starts, and how much a cheaper canvas mask could save). Medians over --iters runs, GPU-synchronized. """ from __future__ import annotations @@ -37,11 +37,16 @@ def main(argv=None): ap.add_argument("--model", default="trillionlabs/Trida2.0-4B") ap.add_argument("--gen-block", type=int, default=None, help="N (canvas 2N-1); default 4") ap.add_argument("--iters", type=int, default=20) + ap.add_argument("--context", type=int, default=0, + help="prefix length in tokens (default ~400); agent prompts are 14K-33K") a = ap.parse_args(argv) eng = Engine(a.model, gen_block=a.gen_block, prompt_cache=False) rt, n, mask = eng.rt, eng.n, eng.mask_id blk = 2 * n - 1 prompt = eng.encode(eng.render(LONG_PROMPT, enable_thinking=False)) + if a.context > len(prompt): + body = prompt[3:-6] + prompt = prompt[:3] + (body * (a.context // len(body) + 1))[: a.context - 9] + prompt[-6:] cache = rt.make_cache() rt.prefill(cache, prompt) mx.eval([l.ssm if isinstance(l, GDNState) else l.keys for l in cache.layers]) @@ -94,6 +99,14 @@ def run_canvas(rows_=None, adv=n): t_cv_unfused = timeit(run_canvas, a.iters) rt.fused_gdn = True t_cv = timeit(run_canvas, a.iters) + # Upper bounds for a cheaper canvas mask (both give wrong logits; timing only): + # no mask at all, and MLX's built-in "causal" mask (no mask array read by the kernel). + real_mask = rt._canvas_mask + rt._canvas_mask = lambda *_: None + t_cv_nomask = timeit(run_canvas, a.iters) + rt._canvas_mask = lambda *_: "causal" + t_cv_causal = timeit(run_canvas, a.iters) + rt._canvas_mask = real_mask t_cold = timeit(lambda: run_canvas(rows_=n, adv=1), a.iters) t_cold_all = timeit(lambda: run_canvas(rows_=None, adv=1), a.iters) t_rej = timeit(lambda: run_canvas(adv=2), a.iters) @@ -125,7 +138,7 @@ def attn(): if not l.is_linear: kv = cache.layers[i] m = rt._canvas_mask(kv.offset, L, n) if L > 1 else None - outs.append(rt._attention(l.self_attn, x, kv, m)) + outs.append(rt._attention(l.self_attn, x, kv, m, kv.offset)) kv.trim(L) return outs t_attn = timeit(attn, a.iters) @@ -154,6 +167,8 @@ def gdn(): print(f"\ncanvas variants:") print(f" verify step, fused GDN {t_cv:7.2f} ms") print(f" verify step, unfused GDN {t_cv_unfused:7.2f} ms (GDN component {c7u[3]:.2f} ms vs fused {c7[3]:.2f} ms)") + print(f" verify step, no attention mask {t_cv_nomask:7.2f} ms (bound; wrong output)") + print(f" verify step, built-in causal mask {t_cv_causal:7.2f} ms (bound; wrong output)") print(f" cold step, LM head rows 0..{n - 1} {t_cold:7.2f} ms") print(f" cold step, LM head all rows {t_cold_all:7.2f} ms") print(f" verify + reject (commit adv=2) {t_rej:7.2f} ms (commit recompute cost {t_rej - t_cv:+.2f} ms)") From 5a6aa21befeaba71dfab391d3874a29c30850110 Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Tue, 6 Oct 2026 08:37:36 +0000 Subject: [PATCH 7/9] docs(inference/mlx): correct mlx-vlm attribution in vision.py header mlx-vlm is by Prince Canuma (Blaizzy), not Apple; matches NOTICE. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- inference/mlx/trida_mlx/vision.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/inference/mlx/trida_mlx/vision.py b/inference/mlx/trida_mlx/vision.py index 3e2e91b..839a06c 100644 --- a/inference/mlx/trida_mlx/vision.py +++ b/inference/mlx/trida_mlx/vision.py @@ -8,8 +8,8 @@ * ``VisionTower`` the Qwen3-VL vision encoder: Conv3d patch embed, bilinearly interpolated learned position embedding, 2-D rotary ViT blocks and the 2x2 patch merger/projector into the LM hidden size. - Adapted from mlx-vlm's ``qwen3_vl/vision.py`` (MIT, Apple Inc. and - contributors) so the parameter names match HF/mlx-vlm checkpoints. + Adapted from mlx-vlm's ``qwen3_vl/vision.py`` (MIT, © 2025 Prince Canuma) + so the parameter names match HF/mlx-vlm checkpoints. * ``mrope_*`` interleaved multimodal RoPE positions for prompts that contain images (text tokens get equal t/h/w positions, i.e. ordinary RoPE). """ From ab3e499ff18765b227dd04c1bbbcdf67dc5e1b2b Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Tue, 6 Oct 2026 08:39:12 +0000 Subject: [PATCH 8/9] test(inference/mlx): restore offline tests and add image-input tests Brings back inference/mlx/tests (reverts d344fdc): tiny random Qwen3.5-hybrid decoder invariants (test_tiny.py), tiny raw-HF checkpoint end to end (test_e2e.py, updated for load_model's vision return and _normalize_messages' image list), the chat-template fixture with its NOTICE / COMPLIANCE entries, the pytest dev group and config, and the README Tests section. Adds test_vision.py on a tiny multimodal checkpoint (tiny_vl.py): image expansion, greedy self-spec = greedy AR with an image, chunked prefill across an image, MRoPE position advance, prompt cache never reusing across different images, convert keeping the vision tower bit-identical, and the server image_url path. 27 tests, ~30 s on CPU (uv run pytest). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- CHANGELOG.md | 2 +- COMPLIANCE.md | 1 + NOTICE | 7 + inference/mlx/README.md | 18 ++ inference/mlx/pyproject.toml | 9 +- .../tests/fixtures/trida_chat_template.jinja | 154 +++++++++ inference/mlx/tests/test_e2e.py | 229 ++++++++++++++ inference/mlx/tests/test_tiny.py | 297 ++++++++++++++++++ inference/mlx/tests/test_vision.py | 111 +++++++ inference/mlx/tests/tiny_ckpt.py | 96 ++++++ inference/mlx/tests/tiny_vl.py | 62 ++++ inference/mlx/uv.lock | 43 +++ 12 files changed, 1027 insertions(+), 2 deletions(-) create mode 100644 inference/mlx/tests/fixtures/trida_chat_template.jinja create mode 100644 inference/mlx/tests/test_e2e.py create mode 100644 inference/mlx/tests/test_tiny.py create mode 100644 inference/mlx/tests/test_vision.py create mode 100644 inference/mlx/tests/tiny_ckpt.py create mode 100644 inference/mlx/tests/tiny_vl.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 32aa065..c7744d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `inference/mlx/`: on-device (Apple Silicon, MLX) serving of Trida2.0-4B with lossless self-speculative decoding (`bd_bidir_shift` b7/g4 semantics), an OpenAI-compatible server with - tool calls and a prompt cache, a minimal agent loop, and MLX quantization. + tool calls and a prompt cache, a minimal agent loop, MLX quantization, and offline tests. - `inference/mlx/`: image input for vision-language checkpoints (Trida-2.0-4B-1006): Qwen3.5 vision encoder + projector, HF-identical preprocessing, multimodal RoPE, image-aware prompt cache, OpenAI `image_url` parts in the server. diff --git a/COMPLIANCE.md b/COMPLIANCE.md index f6cd428..5e0b86c 100644 --- a/COMPLIANCE.md +++ b/COMPLIANCE.md @@ -13,6 +13,7 @@ model-card license. | IFEval scorer | `inference/ifeval_lib/` | [google-research · instruction_following_eval](https://github.com/google-research/google-research/tree/master/instruction_following_eval) | **Apache-2.0** | | mlx-vlm Qwen3-VL vision encoder (adapted) | `inference/mlx/trida_mlx/vision.py` | [Blaizzy/mlx-vlm](https://github.com/Blaizzy/mlx-vlm) `mlx_vlm/models/qwen3_vl/vision.py` | **MIT** | | MLX gated-delta Metal kernel (adapted) | `inference/mlx/trida_mlx/kernels.py` | [ml-explore/mlx-lm](https://github.com/ml-explore/mlx-lm) `mlx_lm/models/gated_delta.py` | **MIT** | +| Qwen3.5 / Trida chat template (test fixture) | `inference/mlx/tests/fixtures/trida_chat_template.jinja` | shipped with `trillionlabs/Trida2.0-4B` (from Qwen3.5) | **Apache-2.0** | These are the only third-party source code redistributed in the repository (see [`NOTICE`](NOTICE)). Everything else below is installed or fetched separately. diff --git a/NOTICE b/NOTICE index 9f5f6dc..3a61454 100644 --- a/NOTICE +++ b/NOTICE @@ -47,6 +47,13 @@ mlx-vlm Qwen3-VL / Qwen3.5 vision encoder position embedding, 2-D rotary blocks, patch merger) is adapted from mlx-vlm so parameter names match its checkpoints. See its header. +Qwen3.5 chat template (test fixture) + Copyright Alibaba Cloud (Qwen team). + License: Apache License 2.0 + Bundled: inference/mlx/tests/fixtures/trida_chat_template.jinja — the chat template + shipped with trillionlabs/Trida2.0-4B (inherited from Qwen3.5), unmodified, + used only by the offline tests. + The components above are the third-party source code redistributed in this repository. They do not share a single license: see COMPLIANCE.md section 1. diff --git a/inference/mlx/README.md b/inference/mlx/README.md index d4e48bf..94b3908 100644 --- a/inference/mlx/README.md +++ b/inference/mlx/README.md @@ -153,6 +153,24 @@ system block, and one just before the last `<|im_start|>` of the previous prompt turn's re-rendered history still agrees with it). A follow-up turn then only prefills the new messages; `usage.prompt_tokens_details.cached_tokens` shows the reuse. +## Tests + +```bash +uv run pytest # all offline tests (~30 s) +uv run python tests/test_tiny.py # decoder invariants on a tiny random Qwen3.5-hybrid (CPU or GPU) +uv run python tests/test_e2e.py # tiny raw-HF checkpoint: loader, convert q8, prompt cache, server, parsing +uv run pytest tests/test_vision.py # tiny multimodal checkpoint: image input end to end +``` + +Pinned: chunked prefill = one-shot; canvas clean rows = AR logits; MASK rows read the block-end GDN +state (independent per-token re-derivation); `commit(adv)` = AR state for every `adv`; greedy +self-spec = greedy AR (also with oracle drafts that exercise the accept path, ~3 tok/forward); +top-k/top-p matches the reference sampler; sampled self-spec reproduces the AR sampling +distribution; snapshot/restore = fresh prefill; server non-stream / SSE / completions round trip; +with an image: greedy self-spec = greedy AR, chunked prefill = one-shot, MRoPE positions, no prompt-cache +reuse across different images, convert keeps the vision tower bit-identical, server `image_url` round trip. +On a Mac the same tests run on the GPU and therefore also cover the Metal kernel path. + ## Notes and limits - **bf16/quantized numerics.** The canvas runs 7-row matmuls and the AR step 1-row matmuls, which diff --git a/inference/mlx/pyproject.toml b/inference/mlx/pyproject.toml index 0c6f636..0258c95 100644 --- a/inference/mlx/pyproject.toml +++ b/inference/mlx/pyproject.toml @@ -11,7 +11,7 @@ requires-python = ">=3.10,<3.14" license = { text = "Apache-2.0" } authors = [{ name = "Trillion Labs" }] dependencies = [ - # Metal build on macOS; the CPU build elsewhere (Linux, no GPU). + # Metal build on macOS; the CPU build elsewhere (Linux CI / offline tests only). "mlx>=0.29; sys_platform == 'darwin'", "mlx[cpu]>=0.29; sys_platform == 'linux'", "mlx-lm>=0.31", # Qwen3.5 gated-delta model + Metal kernel we build on @@ -28,9 +28,16 @@ trida-mlx-verify = "trida_mlx.verify:main" trida-mlx-bench = "trida_mlx.bench:main" trida-mlx-profile = "trida_mlx.profile:main" +[dependency-groups] +dev = ["pytest>=8"] + [tool.setuptools] packages = ["trida_mlx"] [tool.uv] # resolve only for the platforms this package runs on environments = ["sys_platform == 'darwin' and platform_machine == 'arm64'", "sys_platform == 'linux'"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = [".", "tests"] diff --git a/inference/mlx/tests/fixtures/trida_chat_template.jinja b/inference/mlx/tests/fixtures/trida_chat_template.jinja new file mode 100644 index 0000000..ade1daa --- /dev/null +++ b/inference/mlx/tests/fixtures/trida_chat_template.jinja @@ -0,0 +1,154 @@ +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {%- if content %} + {{- '\n\n' + content }} + {%- endif %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if ns.multi_step_tool %} + {{- raise_exception('No user query found in messages.') }} +{%- endif %} +{%- for message in messages %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "system" %} + {%- if not loop.first %} + {{- raise_exception('System message must be at the beginning.') }} + {%- endif %} + {%- elif message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {%- if (loop.index0 > ns.last_query_index) or (keep_all_reasoning is defined and keep_all_reasoning) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- else %} + {{- raise_exception('Unexpected message role.') }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} diff --git a/inference/mlx/tests/test_e2e.py b/inference/mlx/tests/test_e2e.py new file mode 100644 index 0000000..91f8db9 --- /dev/null +++ b/inference/mlx/tests/test_e2e.py @@ -0,0 +1,229 @@ +"""End-to-end (offline) checks on a tiny raw-HF checkpoint: loader, convert (q8), engine +prompt cache, OpenAI server (non-stream + SSE + tools), and output parsing.""" +import json +import os +import socket +import sys +import tempfile +import threading +import urllib.error +import urllib.request +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import mlx.core as mx + +import tiny_ckpt +from test_tiny import tiny_runtime +from trida_mlx import convert +from trida_mlx.decode import DecodeStats, SamplingParams +from trida_mlx.engine import Engine +from trida_mlx.model import TridaRuntime, load_model +from trida_mlx.parsing import StreamSplitter, parse_tool_calls, split_reasoning + +TMP = Path(tempfile.mkdtemp(prefix="trida_e2e_")) +CKPT, MASK_ID = tiny_ckpt.build(TMP / "raw") + + +def test_raw_hf_loader_matches_reference(): + model, _v, cfg, _ = load_model(str(CKPT), dtype=mx.float32) + rt = TridaRuntime(model) + ref = tiny_runtime(seed=0) + p = [5, 17, 99, 3, 250, 8, 41] + a = rt.prefill(rt.make_cache(), p) + b = ref.prefill(ref.make_cache(), p) + # bf16 round trip of the weights -> small differences only + assert float(mx.max(mx.abs(a - b))) < 0.25 * float(mx.max(mx.abs(b))), (a[:5], b[:5]) + assert int(mx.argmax(a)) == int(mx.argmax(b)) + + +def test_convert_q8_and_selfspec_lossless_on_quantized(): + out = TMP / "q8" + convert.main(["--model", str(CKPT), "--out", str(out), "--bits", "8", "--group-size", "32"]) + eng = Engine(str(out)) + assert eng.mask_id == MASK_ID + p = eng.encode(eng.render([{"role": "user", "content": "hello world"}])) + sp = SamplingParams(temperature=0.0) + a = [t for c in eng.generate_ids(p, max_tokens=80, sampling=sp, mode="causal") for t in c] + eng.reset_cache() + b = [t for c in eng.generate_ids(p, max_tokens=80, sampling=sp, mode="self-spec") for t in c] + assert a == b + + +def test_prompt_cache_reuse(): + eng = Engine(str(CKPT)) + sp = SamplingParams(temperature=0.0) + msgs = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello world"}] + p1 = eng.encode(eng.render(msgs)) + out1 = [t for c in eng.generate_ids(p1, max_tokens=30, sampling=sp) for t in c] + msgs2 = msgs + [{"role": "assistant", "content": eng.tokenizer.decode(out1)}, + {"role": "user", "content": "and the weather?"}] + p2 = eng.encode(eng.render(msgs2)) + st = DecodeStats() + out2 = [t for c in eng.generate_ids(p2, max_tokens=30, sampling=sp, stats=st) for t in c] + assert st.reused_tokens > len(p1) // 2, st.as_dict() + # same result as a cold engine + eng2 = Engine(str(CKPT), prompt_cache=False) + ref2 = [t for c in eng2.generate_ids(p2, max_tokens=30, sampling=sp) for t in c] + assert out2 == ref2 + + +def test_parsing(): + c, calls = parse_tool_calls('ok {"name": "f", "arguments": {"x": 1}}}') + assert c == "ok" and calls == [{"name": "f", "arguments": {"x": 1}}] + _, calls = parse_tool_calls("\n\n\nSeoul\n\n\n") + assert calls == [{"name": "g", "arguments": {"city": "Seoul"}}] + assert split_reasoning("abc\n\nhi", True) == ("abc", "hi") + s = StreamSplitter(True, parse_tools=True) + text = 'let me think\n\nSure.{"name": "f", "arguments": {}}' + parts = [] + for ch in text: + parts += s.feed(ch) + rest, calls = s.finish() + parts += rest + reasoning = "".join(t for k, t in parts if k == "reasoning") + content = "".join(t for k, t in parts if k == "content") + assert reasoning == "let me think" and content == "Sure." and calls[0]["name"] == "f", (parts, calls) + + +def test_normalize_messages_against_real_trida_template(): + """Hermes-style histories must render with the real Trida chat template.""" + from transformers import AutoTokenizer + from trida_mlx.server import _normalize_messages + tok = AutoTokenizer.from_pretrained(str(CKPT)) + tok.chat_template = (Path(__file__).parent / "fixtures" / "trida_chat_template.jinja").read_text() + tools = [{"type": "function", "function": {"name": "terminal", "description": "run", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}}}}] + msgs = [ + {"role": "system", "content": "You are Hermes."}, + {"role": "developer", "content": "Be terse."}, + {"role": "user", "content": [{"type": "text", "text": "list files"}]}, + {"role": "assistant", "content": None, "reasoning": "use terminal", + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "terminal", "arguments": '{"command": "ls"}'}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "a.py\nb.py"}, + {"role": "system", "content": "Context was compacted."}, + {"role": "user", "content": "thanks, now count them"}, + ] + norm, _imgs = _normalize_messages(msgs) + assert [m["role"] for m in norm] == ["system", "user", "assistant", "tool", "user", "user"] + text = tok.apply_chat_template(norm, tools=tools, add_generation_prompt=True, tokenize=False, enable_thinking=True) + assert "You are Hermes.\n\nBe terse." in text + assert "\n\nls\n" in text + assert "[System note]\nContext was compacted." in text and text.endswith("\n") + + +def test_tool_param_coercion(): + tools = [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": { + "path": {"type": "string"}, "n": {"type": "integer"}, "flags": {"type": "array"}}}}}] + _, calls = parse_tool_calls("\n\n\n007\n\n" + "\n5\n\n\n[\"-a\"]\n\n" + "\n", tools) + assert calls == [{"name": "f", "arguments": {"path": "007", "n": 5, "flags": ["-a"]}}], calls + + +def test_multi_slot_cache_survives_side_requests(): + eng = Engine(str(CKPT), cache_slots=3) + sp = SamplingParams(temperature=0.0) + run = lambda msgs, st=None: [t for c in eng.generate_ids(eng.encode(eng.render(msgs)), max_tokens=12, + sampling=sp, stats=st) for t in c] + main = [{"role": "system", "content": "main agent " * 30}, {"role": "user", "content": "hello world"}] + out = run(main) + run([{"role": "user", "content": "write a title for: hello"}]) # side request (title) + run([{"role": "user", "content": "summarize: the weather in seoul"}]) # another side request + st = DecodeStats() + run(main + [{"role": "assistant", "content": eng.tokenizer.decode(out)}, + {"role": "user", "content": "and then?"}], st) + assert st.reused_tokens > 100, st.as_dict() + assert len(eng.slots) == 3 + + +def test_system_block_snapshot_reused_when_user_turn_changes(): + """Hermes rewrites the last user message on a truncation retry; the (long) system block + must still be reused.""" + eng = Engine(str(CKPT), cache_slots=1) + sp = SamplingParams(temperature=0.0) + sys_msg = {"role": "system", "content": "tool schemas " * 60} + p1 = eng.encode(eng.render([sys_msg, {"role": "user", "content": "list the files"}])) + for _ in eng.generate_ids(p1, max_tokens=5, sampling=sp): + pass + st = DecodeStats() + p2 = eng.encode(eng.render([sys_msg, {"role": "user", "content": "list the files [System: continue]"}])) + out = [t for c in eng.generate_ids(p2, max_tokens=5, sampling=sp, stats=st) for t in c] + n_sys = p1.index(eng.im_start, 1) # tokens up to the end of the system block + assert st.reused_tokens == n_sys, (st.reused_tokens, n_sys) + ref = Engine(str(CKPT), prompt_cache=False) + assert out == [t for c in ref.generate_ids(p2, max_tokens=5, sampling=sp) for t in c] + + +def _free_port(): + s = socket.socket(); s.bind(("127.0.0.1", 0)); p = s.getsockname()[1]; s.close(); return p + + +def test_server_roundtrip(): + from trida_mlx import server + port = _free_port() + t = threading.Thread(target=server.main, args=(["--model", str(CKPT), "--port", str(port), + "--max-tokens", "40", "--prefill-chunk", "16", + "--context-length", "2048"],), daemon=True) + t.start() + base = f"http://127.0.0.1:{port}/v1" + import time + for _ in range(200): + try: + urllib.request.urlopen(f"{base}/models", timeout=1); break + except Exception: + time.sleep(0.1) + tools = [{"type": "function", "function": {"name": "now", "description": "time", + "parameters": {"type": "object", "properties": {}}}}] + body = {"messages": [{"role": "user", "content": "hello"}], "tools": tools, "temperature": 0} + r = json.loads(urllib.request.urlopen(urllib.request.Request( + f"{base}/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"})).read()) + assert r["choices"][0]["message"]["role"] == "assistant" + assert r["usage"]["trida"]["new_tokens"] > 0 + body["stream"] = True + lines = urllib.request.urlopen(urllib.request.Request( + f"{base}/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"})).read().decode().splitlines() + datas = [l[6:] for l in lines if l.startswith("data: ")] + assert datas[-1] == "[DONE]" + usage = json.loads(datas[-2])["usage"] + assert usage["prompt_tokens_details"]["cached_tokens"] > 0 # second identical request hits the cache + assert any(l.startswith(": prefill") for l in lines) or usage["prompt_tokens_details"]["cached_tokens"] > 0 + r = json.loads(urllib.request.urlopen(urllib.request.Request( + f"{base}/completions", data=json.dumps({"prompt": "hello", "max_tokens": 10}).encode(), + headers={"Content-Type": "application/json"})).read()) + assert r["usage"]["completion_tokens"] <= 10 + models = json.loads(urllib.request.urlopen(f"{base}/models").read()) + assert models["data"][0]["context_length"] == 2048 and models["data"][0]["id"] == "trida2.0-4b" + # context overflow -> OpenAI-style 400 that harnesses recognise (and compress) + big = {"messages": [{"role": "user", "content": "hello world " * 3000}], "stream": True} + try: + urllib.request.urlopen(urllib.request.Request(f"{base}/chat/completions", data=json.dumps(big).encode(), + headers={"Content-Type": "application/json"})) + assert False, "expected 400" + except urllib.error.HTTPError as e: + err = json.loads(e.read())["error"] + assert e.code == 400 and err["code"] == "context_length_exceeded" and "maximum context length" in err["message"] + # keep-alive comments during a long (uncached) prefill + body2 = {"messages": [{"role": "user", "content": "the weather in seoul is sunny " * 40}], "stream": True, + "max_tokens": 4} + lines = urllib.request.urlopen(urllib.request.Request( + f"{base}/chat/completions", data=json.dumps(body2).encode(), + headers={"Content-Type": "application/json"})).read().decode().splitlines() + assert any(l.startswith(": prefill") for l in lines), lines[:5] + + +if __name__ == "__main__": + fails = 0 + for name, fn in list(globals().items()): + if name.startswith("test_"): + try: + fn(); print(f" ok {name}") + except Exception as e: # noqa: BLE001 + import traceback; traceback.print_exc() + fails += 1; print(f" FAIL {name}: {type(e).__name__}: {e}") + sys.stdout.flush() + os._exit(1 if fails else 0) diff --git a/inference/mlx/tests/test_tiny.py b/inference/mlx/tests/test_tiny.py new file mode 100644 index 0000000..7c9f01e --- /dev/null +++ b/inference/mlx/tests/test_tiny.py @@ -0,0 +1,297 @@ +"""Correctness tests on a tiny random Qwen3.5-hybrid model (CPU is fine; no weights needed). + + python -m pytest inference/mlx/tests -q # or: python inference/mlx/tests/test_tiny.py + +What is pinned: + * chunked prefill == one-shot prefill + * canvas clean rows (0..N-1) reproduce the AR logits exactly (the lossless property) + * canvas MASK rows read the *block-end* GDN state (causal_mode=2 semantics) + * commit(adv) leaves the cache identical to having AR-decoded those adv tokens + * greedy self-spec output == greedy AR output, token for token, over long generations + * a snapshot/restore round trip (prompt cache) reproduces a fresh prefill +""" +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from trida_mlx.decode import DecodeStats, SamplingParams, ar_generate, selfspec_generate +from trida_mlx.model import GDNState, TridaRuntime, build_model +from mlx_lm.models.gated_delta import compute_g + +MASK = 299 +TINY = dict( + model_type="qwen3_5_text", hidden_size=64, intermediate_size=128, num_hidden_layers=8, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, rms_norm_eps=1e-6, vocab_size=300, + linear_num_value_heads=4, linear_num_key_heads=2, linear_key_head_dim=32, linear_value_head_dim=16, + linear_conv_kernel_dim=4, tie_word_embeddings=True, full_attention_interval=4, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 0.25, + "mrope_section": [1, 1, 0]}, +) + + +def tiny_runtime(seed=0, sharp=4.0): + mx.random.seed(seed) + model = build_model(TINY) + # sharpen the output distribution so greedy decoding has structure (and fewer near-ties) + emb = model.language_model.model.embed_tokens + emb.weight = emb.weight * sharp + mx.eval(model.parameters()) + return TridaRuntime(model, prefill_chunk=512) + + +def rand_prompt(n, seed=1): + rng = np.random.default_rng(seed) + return rng.integers(0, 290, size=n).tolist() + + +def close(a, b, tol=1e-4): + return float(mx.max(mx.abs(a - b)).item()) < tol + + +def test_chunked_prefill(): + rt = tiny_runtime() + p = rand_prompt(37) + c1 = rt.make_cache(); l1 = rt.prefill(c1, p) + rt2 = tiny_runtime(); rt2.prefill_chunk = 5 + c2 = rt2.make_cache(); l2 = rt2.prefill(c2, p) + assert close(l1, l2) + + +def _ar_logits(rt, prompt, cont): + c = rt.make_cache() + out = [rt.prefill(c, prompt)] + for t in cont: + out.append(rt.ar_step(c, t)) + return out, c + + +def test_canvas_clean_rows_are_ar_logits(): + rt = tiny_runtime(); n = 4 + prompt = rand_prompt(20) + cont = rand_prompt(n, seed=7) # p, s0, s1, s2 + ar, _ = _ar_logits(rt, prompt, cont) + c = rt.make_cache(); rt.prefill(c, prompt) + L = rt.canvas(c, cont + [MASK] * (n - 1), n) + for i in range(n): # row i == AR logits after cont[:i+1] + assert close(L[i], ar[i + 1]), f"row {i}" + + +def test_commit_matches_ar_state(): + n = 4 + for adv in range(1, n + 1): + rt = tiny_runtime() + prompt = rand_prompt(15) + cont = rand_prompt(n, seed=3) + c = rt.make_cache(); rt.prefill(c, prompt) + rt.canvas(c, cont + [MASK] * (n - 1), n) + rt.commit(c, adv) + nxt = 42 + got = rt.ar_step(c, nxt) + ar, _ = _ar_logits(rt, prompt, cont[:adv] + [nxt]) + assert close(got, ar[-1]), f"adv={adv}" + assert c.tokens == prompt + cont[:adv] + [nxt] + + +def test_mask_rows_read_block_end_state(): + """Independent re-derivation of one GDN layer's canvas output with a per-token loop.""" + rt = tiny_runtime(); n = 4 + layer = next(l for l in rt.backbone.layers if l.is_linear) + lin = layer.linear_attn + x = mx.random.normal((1, 2 * n - 1, TINY["hidden_size"])) + st = GDNState(conv=mx.random.normal((1, 3, lin.conv_dim)), + ssm=mx.random.normal((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim)) * 0.1) + out, _, _ = rt._gdn(lin, x, st, n) + # manual: conv -> q,k,v -> step recurrence, record states + qkv = lin.in_proj_qkv(x); z = lin.in_proj_z(x).reshape(1, 2 * n - 1, lin.num_v_heads, lin.head_v_dim) + ci = mx.concatenate([st.conv, qkv], axis=1) + co = nn.silu(lin.conv1d(ci)) + q, k, v = mx.split(co, [lin.key_dim, 2 * lin.key_dim], -1) + T = 2 * n - 1 + q = q.reshape(1, T, lin.num_k_heads, -1); k = k.reshape(1, T, lin.num_k_heads, -1) + v = v.reshape(1, T, lin.num_v_heads, -1) + s = k.shape[-1] ** -0.5 + q = s * s * mx.fast.rms_norm(q, None, 1e-6); k = s * mx.fast.rms_norm(k, None, 1e-6) + rep = lin.num_v_heads // lin.num_k_heads + q = mx.repeat(q, rep, 2); k = mx.repeat(k, rep, 2) + beta = mx.sigmoid(lin.in_proj_b(x)); g = compute_g(lin.A_log, lin.in_proj_a(x), lin.dt_bias) + S = st.ssm; ys = [] + for t in range(T): + S = S * g[:, t][..., None, None] + kv = (S * k[:, t][:, :, None, :]).sum(-1) + S = S + k[:, t][:, :, None, :] * ((v[:, t] - kv) * beta[:, t][..., None])[..., None] + ys.append((S * q[:, t][:, :, None, :]).sum(-1)) + S_end = S + for t in range(n, T): + ys[t] = (S_end * q[:, t][:, :, None, :]).sum(-1) + y = mx.stack(ys, axis=1) + ref = lin.out_proj(lin.norm(y, z).reshape(1, T, -1)) + assert close(out, ref, 1e-4) + + +def test_fused_gdn_matches_unfused(): + """Metal-only: the fused canvas GDN kernel == two mlx-lm kernel calls + block-end readout.""" + from trida_mlx.kernels import fused_available + if not fused_available(): + print(" (skipped: no Metal)") + return + rt = tiny_runtime(); n = 4 + lin = next(l for l in rt.backbone.layers if l.is_linear).linear_attn + x = mx.random.normal((1, 2 * n - 1, TINY["hidden_size"])) + st = GDNState(conv=mx.random.normal((1, 3, lin.conv_dim)), + ssm=mx.random.normal((1, lin.num_v_heads, lin.head_v_dim, lin.head_k_dim)) * 0.1) + rt.fused_gdn = True + o1, _, r1 = rt._gdn(lin, x, st, n) + rt.fused_gdn = False + o2, _, r2 = rt._gdn(lin, x, st, n) + assert close(o1, o2, 1e-5) and close(r1.s_clean, r2.s_clean, 1e-6) + # clean rows must be bit-identical to the unfused (= AR) kernel + assert bool(mx.array_equal(o1[:, :n], o2[:, :n]).item()) + + +def _collect(gen): + out = [] + for chunk in gen: + out.extend(chunk) + return out + + +def test_greedy_selfspec_is_lossless(): + sp = SamplingParams(temperature=0.0) + for seed in range(3): + rt = tiny_runtime(seed=seed) + prompt = rand_prompt(25, seed=seed + 10) + ar = _collect(ar_generate(rt, rt.make_cache(), prompt, max_new_tokens=120, eos_ids=set(), sp=sp)) + from trida_mlx.decode import DecodeStats + st = DecodeStats() + ss = _collect(selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=120, eos_ids=set(), + sp=sp, mask_id=MASK, n=4, stats=st)) + assert ss == ar, f"seed {seed}: first diff at {next(i for i,(a,b) in enumerate(zip(ss,ar)) if a!=b)}" + assert st.forwards <= len(ar) + + +def test_greedy_lossless_with_oracle_drafts(): + """Random models rarely accept drafts, so also drive the accept path: patch the canvas + so MASK rows 'predict' the true AR continuation and check output + acceptance.""" + sp = SamplingParams(temperature=0.0) + rt = tiny_runtime(seed=5) + prompt = rand_prompt(18, seed=4) + ar = _collect(ar_generate(rt, rt.make_cache(), prompt, max_new_tokens=60, eos_ids=set(), sp=sp)) + full = prompt + ar + orig = rt.canvas + + # replace the MASK draft rows with a one-hot on the true AR continuation + def patched(cache, toks, n, rows=None): + base = cache.length + L = orig(cache, toks, n, rows=rows) + rows = [] + V = L.shape[-1] + for i in range(L.shape[0]): + want = base + i + 1 + if i >= 1 and want < len(full) and (toks[i] == MASK): + rows.append(mx.where(mx.arange(V) == full[want], 100.0, 0.0)) + else: + rows.append(L[i]) + return mx.stack(rows) + + rt.canvas = patched + from trida_mlx.decode import DecodeStats + st = DecodeStats() + ss = _collect(selfspec_generate(rt, rt.make_cache(), prompt, max_new_tokens=60, eos_ids=set(), + sp=sp, mask_id=MASK, n=4, stats=st)) + assert ss == ar + assert st.accept_hist[-1] > 0 and st.tokens_per_forward > 2.5, st.as_dict() + + +def test_snapshot_restore(): + rt = tiny_runtime() + p = rand_prompt(30) + c = rt.make_cache(); rt.prefill(c, p[:20]) + snap = c.snapshot() + rt.prefill(c, p[20:]) + for t in rand_prompt(9, seed=9): + rt.ar_step(c, t) + c.restore(snap) + l1 = rt.prefill(c, p[20:]) + l2 = rt.prefill(rt.make_cache(), p) + assert close(l1, l2) + + +def test_sampling_runs(): + sp = SamplingParams(temperature=1.0, top_k=20, top_p=0.9, seed=0) + rt = tiny_runtime() + out = _collect(selfspec_generate(rt, rt.make_cache(), rand_prompt(12), max_new_tokens=40, + eos_ids=set(), sp=sp, mask_id=MASK, n=4)) + assert len(out) == 40 + + +def test_truncated_probs_matches_reference(): + from trida_mlx.decode import truncated_probs + + def ref(x, t, k, p): # numpy port of the SGLang reference _sampling_probs + x = x / t + kth = np.sort(x, axis=-1)[:, -k][:, None] + x = np.where(x < kth, -np.inf, x) + order = np.argsort(-x, axis=-1, kind="stable") + sx = np.take_along_axis(x, order, -1) + e = np.exp(sx - sx.max(-1, keepdims=True)); pr = e / e.sum(-1, keepdims=True) + sx = np.where(np.cumsum(pr, -1) - pr >= p, -np.inf, sx) + out = np.full_like(x, -np.inf); np.put_along_axis(out, order, sx, -1) + e = np.exp(out - out.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True) + + rng = np.random.default_rng(0) + x = rng.normal(size=(7, 1000)).astype(np.float32) * 3 + for t, k, p in [(1.0, 50, 0.95), (0.7, 20, 0.8), (1.0, 5, 1.0)]: + got = np.array(truncated_probs(mx.array(x), SamplingParams(temperature=t, top_k=k, top_p=p))) + assert np.allclose(got, ref(x, t, k, p), atol=1e-5), (t, k, p) + + +def test_speculative_sampling_preserves_distribution(): + """Sampled self-spec must reproduce the AR sampling distribution. Random models almost + never accept drafts, so the MASK rows are patched to reuse a neighbouring row's logits: + any draft q is valid for speculative sampling, and this one gets accepted often.""" + from collections import Counter + rt = tiny_runtime(seed=2, sharp=1.0) + prompt = rand_prompt(10, seed=2) + sp = SamplingParams(temperature=1.0, top_k=4, top_p=0.9) + orig = rt.canvas + + def patched(cache, toks, n, rows=None): + L = orig(cache, toks, n, rows=rows) + return mx.stack([L[i] if (i == 0 or toks[i] != MASK) else L[i - 1] for i in range(L.shape[0])]) + + def hist(fn, reps, seed0, **kw): + c, st_all = Counter(), DecodeStats() + st_all.accept_hist = [0] * 4 + for i in range(reps): + mx.random.seed(seed0 + i) + st = DecodeStats() + out = _collect(fn(rt, rt.make_cache(), prompt, max_new_tokens=4, eos_ids=set(), sp=sp, stats=st, **kw)) + c[out[3]] += 1 + if st.accept_hist: + st_all.accept_hist = [x + y for x, y in zip(st_all.accept_hist, st.accept_hist)] + return c, st_all + + R = 600 + a, _ = hist(ar_generate, R, 5000) + a2, _ = hist(ar_generate, R, 90000) + rt.canvas = patched + b, st = hist(selfspec_generate, R, 7000, mask_id=MASK, n=4) + tv = lambda x, y: 0.5 * sum(abs(x[k] - y[k]) for k in set(x) | set(y)) / R + assert sum(st.accept_hist[1:]) > R // 4, st.accept_hist # the accept path really ran + assert tv(a, b) < max(0.08, 2.0 * tv(a, a2)), (tv(a, b), tv(a, a2)) + + +if __name__ == "__main__": + fails = 0 + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + try: + fn(); print(f" ok {name}") + except Exception as e: # noqa: BLE001 + fails += 1; print(f" FAIL {name}: {type(e).__name__}: {e}") + sys.exit(1 if fails else 0) diff --git a/inference/mlx/tests/test_vision.py b/inference/mlx/tests/test_vision.py new file mode 100644 index 0000000..7248263 --- /dev/null +++ b/inference/mlx/tests/test_vision.py @@ -0,0 +1,111 @@ +"""Image-input tests on a tiny random multimodal checkpoint (see tiny_vl.py; CPU is fine). + +Pinned: image expansion into per-image keys; greedy self-spec = greedy AR with an image in the +prompt; chunked prefill across an image = one-shot; MRoPE position advance after an image; the +prompt cache never reuses across different images but does reuse past the same one; convert keeps +the vision tower bf16 (bit-identical features); server image_url round trip and 400 on bad data. +""" +import base64, json, os, socket, sys, tempfile, threading, time, urllib.request +from pathlib import Path +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import mlx.core as mx, numpy as np +import tiny_vl +from trida_mlx.engine import Engine +from trida_mlx.decode import SamplingParams, DecodeStats +from trida_mlx import convert + +TMP = Path(tempfile.mkdtemp(prefix="trida_vl_")) +CKPT, IDS = tiny_vl.build(TMP / "raw") +G = SamplingParams(temperature=0.0) + +def gen(eng, prompt, mode="self-spec", n=40, st=None): + return [t for c in eng.generate_ids(prompt, max_tokens=n, sampling=G, mode=mode, stats=st) for t in c] + +def vl_prompt(eng, img_bytes, text="describe this image"): + msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}] + key_n = eng.add_image(img_bytes) + return eng.encode(eng.render(msgs, enable_thinking=False), [key_n]), key_n + +def test_load_and_expand(): + eng = Engine(str(CKPT)) + assert eng.supports_vision + p, (key, n) = vl_prompt(eng, tiny_vl.png(1)) + assert n > 0 and p.count(key) == n + assert IDS["<|image_pad|>"] not in p + +def test_greedy_lossless_with_image(): + eng = Engine(str(CKPT), prompt_cache=False) + p, _ = vl_prompt(eng, tiny_vl.png(2)) + a = gen(eng, p, "causal"); b = gen(eng, p, "self-spec") + assert a == b, (a[:10], b[:10]) + +def test_chunked_prefill_across_image(): + e1 = Engine(str(CKPT), prompt_cache=False) + e2 = Engine(str(CKPT), prompt_cache=False, prefill_chunk=5) + p, _ = vl_prompt(e1, tiny_vl.png(3)); e2.add_image(tiny_vl.png(3)) + assert gen(e1, p, "causal", 20) == gen(e2, p, "causal", 20) + +def test_image_changes_output_and_positions(): + eng = Engine(str(CKPT), prompt_cache=False) + p1, (k1, n1) = vl_prompt(eng, tiny_vl.png(4)) + p2, _ = vl_prompt(eng, tiny_vl.png(5)) + c = eng.rt.make_cache(); l1 = eng.rt.prefill(c, p1) + m = 2; t, h, w = eng.rt.images[k1]["grid"] + assert c.pos == len(p1) - n1 + max(t, h // m, w // m), (c.pos, len(p1), n1) + l2 = eng.rt.prefill(eng.rt.make_cache(), p2) + assert float(mx.max(mx.abs(l1 - l2))) > 1e-3 # different image -> different logits + +def test_prompt_cache_distinguishes_images(): + eng = Engine(str(CKPT), cache_slots=1) + p1, _ = vl_prompt(eng, tiny_vl.png(6)); gen(eng, p1, n=8) + st = DecodeStats() + p2, _ = vl_prompt(eng, tiny_vl.png(7)); out = gen(eng, p2, n=8, st=st) + first_img = next(i for i, t in enumerate(p2) if t < 0) + assert st.reused_tokens <= first_img, (st.reused_tokens, first_img) + ref = Engine(str(CKPT), prompt_cache=False); ref.add_image(tiny_vl.png(7)) + assert out == gen(ref, p2, n=8) + # same image again, follow-up turn -> reuses past the image + st2 = DecodeStats() + msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "describe this image"}]}, + {"role": "assistant", "content": eng.tokenizer.decode(out)}, {"role": "user", "content": "more"}] + k = eng.add_image(tiny_vl.png(7)) + p3 = eng.encode(eng.render(msgs, enable_thinking=False), [k]); gen(eng, p3, n=4, st=st2) + assert st2.reused_tokens > first_img + k[1], (st2.reused_tokens, first_img, k) + +def test_convert_keeps_vision(): + out = TMP / "q8" + convert.main(["--model", str(CKPT), "--out", str(out), "--bits", "8", "--group-size", "32"]) + eng = Engine(str(out), prompt_cache=False) + assert eng.supports_vision and (out / "preprocessor_config.json").exists() + raw = Engine(str(CKPT), prompt_cache=False) + img = tiny_vl.png(8) + ka, _ = eng.add_image(img); kb, _ = raw.add_image(img) + fa = eng.vision(mx.array(eng.rt.images[ka]["pixels"]), [eng.rt.images[ka]["grid"]]) + fb = raw.vision(mx.array(raw.rt.images[kb]["pixels"]), [raw.rt.images[kb]["grid"]]) + assert float(mx.max(mx.abs(fa - fb))) == 0.0 # vision tower kept bf16, bit-identical + p, _ = vl_prompt(eng, img) + assert gen(eng, p, "causal", 20) == gen(eng, p, "self-spec", 20) + +def test_server_image_request(): + from trida_mlx import server + s = socket.socket(); s.bind(("127.0.0.1", 0)); port = s.getsockname()[1]; s.close() + threading.Thread(target=server.main, args=(["--model", str(CKPT), "--port", str(port), "--max-tokens", "12"],), daemon=True).start() + base = f"http://127.0.0.1:{port}/v1" + for _ in range(300): + try: urllib.request.urlopen(f"{base}/models", timeout=1); break + except Exception: time.sleep(0.1) + card = json.loads(urllib.request.urlopen(f"{base}/models").read())["data"][0] + assert card["supports_vision"] is True + url = "data:image/png;base64," + base64.b64encode(tiny_vl.png(9)).decode() + body = {"messages": [{"role": "user", "content": [{"type": "text", "text": "what is in the picture"}, + {"type": "image_url", "image_url": {"url": url}}]}], + "temperature": 0} + r = json.loads(urllib.request.urlopen(urllib.request.Request(f"{base}/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"})).read()) + assert r["usage"]["prompt_tokens"] > 20 and r["usage"]["completion_tokens"] > 0, r["usage"] + bad = {"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}]}]} + try: + urllib.request.urlopen(urllib.request.Request(f"{base}/chat/completions", data=json.dumps(bad).encode(), + headers={"Content-Type": "application/json"})); assert False + except urllib.error.HTTPError as e: + assert e.code == 400 diff --git a/inference/mlx/tests/tiny_ckpt.py b/inference/mlx/tests/tiny_ckpt.py new file mode 100644 index 0000000..31474c0 --- /dev/null +++ b/inference/mlx/tests/tiny_ckpt.py @@ -0,0 +1,96 @@ +"""Build a tiny *raw-HF-format* Trida-like checkpoint (random weights, byte-level BPE +tokenizer, Qwen-style chat template with tools + thinking) for offline end-to-end tests.""" +import json +import os +import sys +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import mlx.core as mx +from mlx.utils import tree_flatten + +from test_tiny import TINY, tiny_runtime # noqa: E402 + +SPECIALS = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "", "", "", + "", "<|mask|>"] + +CHAT_TEMPLATE = r"""{%- if tools %}<|im_start|>system +{%- if messages[0].role == 'system' %} +{{ messages[0].content }} +{%- endif %} +# Tools + +{%- for t in tools %} +{{ t | tojson }} +{%- endfor %} + +Call a tool with {"name": ..., "arguments": {...}}<|im_end|> +{% elif messages[0].role == 'system' %}<|im_start|>system +{{ messages[0].content }}<|im_end|> +{% endif %} +{%- for m in messages %} +{%- if m.role == 'user' %}<|im_start|>user +{{ m.content }}<|im_end|> +{% elif m.role == 'assistant' %}<|im_start|>assistant +{{ m.content }} +{%- if m.tool_calls %}{% for tc in m.tool_calls %}{"name": "{{ tc.function.name }}", "arguments": {{ tc.function.arguments | tojson }}}{% endfor %}{% endif %}<|im_end|> +{% elif m.role == 'tool' %}<|im_start|>user +{{ m.content }}<|im_end|> +{% endif %} +{%- endfor %} +{%- if add_generation_prompt %}<|im_start|>assistant +{%- if enable_thinking is defined and enable_thinking is false %} + + + + +{% else %} + +{% endif %} +{%- endif %}""" + + +def build(out: Path): + from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers + from transformers import PreTrainedTokenizerFast + + out.mkdir(parents=True, exist_ok=True) + tok = Tokenizer(models.BPE()) + tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + tok.decoder = decoders.ByteLevel() + corpus = ["hello world the weather in seoul is sunny", "def fib(n): return n", "tool call json name", + "안녕하세요 온디바이스"] * 20 + trainer = trainers.BpeTrainer(vocab_size=TINY["vocab_size"] - len(SPECIALS), special_tokens=[], + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False) + tok.train_from_iterator(corpus, trainer) + tok.add_special_tokens(SPECIALS) + assert tok.get_vocab_size() <= TINY["vocab_size"], tok.get_vocab_size() + fast = PreTrainedTokenizerFast(tokenizer_object=tok, eos_token="<|im_end|>", pad_token="<|endoftext|>") + fast.chat_template = CHAT_TEMPLATE + fast.save_pretrained(str(out)) + mask_id = fast.convert_tokens_to_ids("<|mask|>") + + rt = tiny_runtime(seed=0) + params = dict(tree_flatten(rt.model.parameters())) + hf = {} + norm_sfx = (".input_layernorm.weight", ".post_attention_layernorm.weight", "model.norm.weight", + ".q_norm.weight", ".k_norm.weight") + for k, v in params.items(): + k2 = k.replace("language_model.model.", "model.").replace("language_model.", "") + if "conv1d.weight" in k2: + v = v.moveaxis(1, 2) # MLX (C, K, 1) -> HF (C, 1, K) + if any(k2.endswith(s) for s in norm_sfx): + v = v - 1.0 # HF Qwen3.5 stores zero-centred RMSNorm weights + hf[k2] = v.astype(mx.bfloat16) if not k2.endswith("A_log") else v + hf["lm_head.weight"] = hf["model.embed_tokens.weight"] # real ckpt ships the tied head too + mx.save_safetensors(str(out / "model.safetensors"), hf) + cfg = dict(TINY, architectures=["Qwen3_5ForCausalLM"]) + (out / "config.json").write_text(json.dumps(cfg, indent=2)) + (out / "block_diffusion.json").write_text(json.dumps({"bd_size": 4, "mask_id": mask_id})) + (out / "generation_config.json").write_text(json.dumps({"eos_token_id": [fast.eos_token_id]})) + return out, mask_id + + +if __name__ == "__main__": + print(build(Path(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tiny_trida"))) diff --git a/inference/mlx/tests/tiny_vl.py b/inference/mlx/tests/tiny_vl.py new file mode 100644 index 0000000..b119cf8 --- /dev/null +++ b/inference/mlx/tests/tiny_vl.py @@ -0,0 +1,62 @@ +"""Tiny Trida-like *multimodal* raw-HF checkpoint (random weights) for offline tests.""" +import json, os, sys +from pathlib import Path +import numpy as np +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import mlx.core as mx +from mlx.utils import tree_flatten +from test_tiny import TINY, tiny_runtime +from trida_mlx.vision import VisionTower, VisionConfig + +SPECIALS = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "", "", "", + "", "<|mask|>", "<|vision_start|>", "<|image_pad|>", "<|vision_end|>", "<|video_pad|>"] +VCFG = dict(depth=2, hidden_size=32, intermediate_size=64, num_heads=2, patch_size=4, temporal_patch_size=2, + spatial_merge_size=2, in_channels=3, out_hidden_size=TINY["hidden_size"], num_position_embeddings=16, + deepstack_visual_indexes=[], model_type="qwen3_5") + +def build(out: Path): + from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers + from transformers import PreTrainedTokenizerFast + out.mkdir(parents=True, exist_ok=True) + tok = Tokenizer(models.BPE()); tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + tok.decoder = decoders.ByteLevel() + corpus = ["hello world the weather in seoul is sunny describe this image", "what is in the picture"] * 20 + tok.train_from_iterator(corpus, trainers.BpeTrainer(vocab_size=TINY["vocab_size"] - len(SPECIALS), special_tokens=[], + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)) + tok.add_special_tokens(SPECIALS) + fast = PreTrainedTokenizerFast(tokenizer_object=tok, eos_token="<|im_end|>", pad_token="<|endoftext|>") + fast.chat_template = (Path(__file__).parent / "fixtures" / "trida_chat_template.jinja").read_text() + fast.save_pretrained(str(out)) + ids = {t: fast.convert_tokens_to_ids(t) for t in SPECIALS} + rt = tiny_runtime(seed=0) + hf = {} + norm_sfx = (".input_layernorm.weight", ".post_attention_layernorm.weight", "model.norm.weight", ".q_norm.weight", ".k_norm.weight") + for k, v in tree_flatten(rt.model.parameters()): + k2 = ("model.language_model." + k[len("language_model.model."):]) if k.startswith("language_model.model.") else k[len("language_model."):] + if "conv1d.weight" in k2: v = v.moveaxis(1, 2) + if any(k2.endswith(s) for s in norm_sfx): v = v - 1.0 + hf[k2] = v if k2.endswith("A_log") else v.astype(mx.bfloat16) + mx.random.seed(7) + vt = VisionTower(VisionConfig.from_dict(VCFG)) + for k, v in tree_flatten(vt.parameters()): + if k == "patch_embed.proj.weight": v = v.transpose(0, 4, 1, 2, 3) # MLX NDHWC -> PyTorch NCDHW + hf["model.visual." + k] = v.astype(mx.bfloat16) + mx.save_safetensors(str(out / "model.safetensors"), hf) + text = dict(TINY); text["model_type"] = "qwen3_5_text" + cfg = {"architectures": ["Qwen3_5ForConditionalGeneration"], "model_type": "qwen3_5", "text_config": text, + "vision_config": VCFG, "tie_word_embeddings": True, "image_token_id": ids["<|image_pad|>"], + "video_token_id": ids["<|video_pad|>"], "vision_start_token_id": ids["<|vision_start|>"], + "vision_end_token_id": ids["<|vision_end|>"]} + (out / "config.json").write_text(json.dumps(cfg, indent=2)) + (out / "preprocessor_config.json").write_text(json.dumps({"patch_size": 4, "temporal_patch_size": 2, "merge_size": 2, + "size": {"shortest_edge": 16 * 16, "longest_edge": 48 * 48}, "image_mean": [0.5] * 3, "image_std": [0.5] * 3})) + (out / "block_diffusion.json").write_text(json.dumps({"bd_size": 4, "mask_id": ids["<|mask|>"]})) + (out / "generation_config.json").write_text(json.dumps({"eos_token_id": [fast.eos_token_id]})) + return out, ids + +def png(seed, w=40, h=28): + from PIL import Image + import io + rng = np.random.default_rng(seed) + img = Image.fromarray(rng.integers(0, 255, (h, w, 3), dtype=np.uint8)) + b = io.BytesIO(); img.save(b, "PNG"); return b.getvalue() diff --git a/inference/mlx/uv.lock b/inference/mlx/uv.lock index b8f5d4d..d9d3c5c 100644 --- a/inference/mlx/uv.lock +++ b/inference/mlx/uv.lock @@ -165,6 +165,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/58/a2/bb081bab032533a855d44de1d56f8e8426114ff1ba5d1f07a438a0a654f8/idna-3.20-py3-none-any.whl", hash = "sha256:ab7ae7122974553370f0bdb919e1a960b2cd1bc1ef0276416d896db81c14582c", size = 69583, upload-time = "2026-09-17T14:11:03.168Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -470,6 +479,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/08/14/f6ba68107680ffa74b39985f3f30884e41318fbc4250caa423c79b4788bb/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7e480451b9fa137494bccd3a7d69adbe8ac65a87d97be61e11f1b1050a5bac3", size = 5860358, upload-time = "2026-07-01T11:56:32.68Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "protobuf" version = "7.36.2" @@ -492,6 +510,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "pyyaml" version = "6.0.3" @@ -742,6 +777,11 @@ dependencies = [ { name = "transformers" }, ] +[package.dev-dependencies] +dev = [ + { name = "pytest" }, +] + [package.metadata] requires-dist = [ { name = "huggingface-hub" }, @@ -752,6 +792,9 @@ requires-dist = [ { name = "transformers", specifier = ">=4.57" }, ] +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=8" }] + [[package]] name = "typer" version = "0.27.2" From 1da50f1718b916a4a9203e137e11e5402d94fcab Mon Sep 17 00:00:00 2001 From: Hyungguk Date: Tue, 13 Oct 2026 05:08:08 +0000 Subject: [PATCH 9/9] chore(inference/mlx): require setuptools>=83 like the other packages Matches the PYSEC-2026-3447 bump main applied to the root and vLLM packages. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013o6LMxkLD6mBHGy1C6MBR7 --- inference/mlx/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inference/mlx/pyproject.toml b/inference/mlx/pyproject.toml index 0258c95..0a6d4cf 100644 --- a/inference/mlx/pyproject.toml +++ b/inference/mlx/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=64"] +requires = ["setuptools>=83"] # >=83: PYSEC-2026-3447 build-backend = "setuptools.build_meta" [project]