| """Multi-turn bounded SP-evict chat in MLX (the 'long-CoT rally' the project wanted). |
| Persistent KV across turns but bounded: each chunk trims cache to the system prompt, prepends |
| 32 SP (compressing the distant conversation) + recent raw window. User msg tokens are force-fed, |
| then the assistant samples its reply. Tests whether reasoning is carried across turns under bounded memory.""" |
| import sys, os, time |
| import numpy as np |
| import mlx.core as mx |
| from mlx_lm.models.cache import make_prompt_cache |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import sp_mlx |
|
|
| import sys as _s |
| RW = int(_s.argv[1]) if len(_s.argv)>1 else 128 |
| C, MAXD, TEMP, MAX_RESP = 64, 4096, 0.6, 500 |
| M = sp_mlx.get() |
| model, tok, pooler, H, eos, embT, mdtype = (M["model"], M["tok"], M["pooler"], M["H"], M["eos"], M["embT"], M["mdtype"]) |
|
|
|
|
| def emb_ids(ids): |
| return mx.zeros((1, 0, H), dtype=mdtype) if len(ids) == 0 else embT(mx.array([ids])) |
|
|
|
|
| def evict(kept): |
| if not MAXD or len(kept) <= MAXD: |
| return kept |
| _, mass = pooler.forward_with_mass(emb_ids(kept).astype(mx.float32)); mx.eval(mass) |
| m = np.array(mass)[0]; idx = np.sort(np.argsort(m)[-MAXD:]) |
| return [kept[i] for i in idx] |
|
|
|
|
| SYSTEM = "You are a helpful math tutor. Reason briefly step by step, then give the final answer." |
| cache = make_prompt_cache(model) |
| sys_ids = tok.encode(SYSTEM) |
| model(mx.array([sys_ids]), cache=cache); MQ = cache[0].offset |
| gen, kept, absorbed = [], [], 0 |
| print(f"SYSTEM seeded MQ={MQ}\n", flush=True) |
|
|
|
|
| def turn(user_msg): |
| global gen, kept, absorbed |
| feed = tok.encode(("<｜end▁of▁sentence｜>" if gen else "") + f"<｜User｜>{user_msg}<｜Assistant｜>", add_special_tokens=False) |
| fi, sampled, start, done = 0, 0, len(gen), False |
| t0 = time.time() |
| while not done: |
| c0 = len(gen); R = min(c0, RW); nd_end = c0 - R |
| if nd_end > absorbed: |
| kept.extend(gen[absorbed:nd_end]); absorbed = nd_end; kept = evict(kept) |
| sp = pooler.forward(emb_ids(kept).astype(mx.float32)).astype(mdtype) |
| parts = [sp] + ([emb_ids(gen[c0 - R:c0])] if R > 0 else []) |
| block = mx.concatenate(parts, axis=1) |
| for c in cache: |
| c.trim(c.offset - MQ) |
| logits = model(mx.zeros((1, block.shape[1]), dtype=mx.int32), cache=cache, input_embeddings=block) |
| last = logits[:, -1, :] |
| for _ in range(C): |
| if fi < len(feed): |
| t = feed[fi]; fi += 1 |
| elif sampled < MAX_RESP: |
| t = int(mx.random.categorical(last * (1.0 / TEMP)).item()); sampled += 1 |
| if t == eos: |
| done = True; break |
| else: |
| done = True; break |
| gen.append(t) |
| logits = model(mx.zeros((1, 1), dtype=mx.int32), cache=cache, input_embeddings=emb_ids([t])) |
| last = logits[:, -1, :] |
| if done: |
| break |
| reply = tok.decode(gen[start + len(feed):]) |
| return reply, time.time() - t0 |
|
|
|
|
| CONV = [ |
| "I have 3 apples and then I buy 5 more. How many apples do I have?", |
| "If I now give away 2 apples, how many do I have left?", |
| "Now I double the number of apples I have. How many is that?", |
| "I eat 4 of them. How many remain?", |
| "Finally, I split the remaining apples equally between 2 friends. How many does each friend get?", |
| ] |
| for i, u in enumerate(CONV, 1): |
| r, dt = turn(u) |
| print(f"--- Turn {i} total_ctx={len(gen)} survivors={len(kept)} {dt:.1f}s ---", flush=True) |
| print(f"USER: {u}", flush=True) |
| print(f"ASSISTANT: {r.strip()[:500]}", flush=True) |
| print("", flush=True) |
| print("MT_MLX_DONE", flush=True) |
|
|
