"""TUI 后端抽象:把界面与 MLX 推理引擎解耦。 界面只依赖 ChatBackend 接口,不 import 任何 MLX 符号,从而能用 FakeBackend 做无模型测试。 load / generate 都是阻塞调用,由 UI 层放到 worker 线程执行。 """ from __future__ import annotations from dataclasses import dataclass, field from typing import Callable, Protocol from collections import OrderedDict @dataclass class GenResult: """一轮生成的汇总。""" text: str # 完整回答(已截断 EOS) n_tokens: int # 新生成 token 数 tok_per_s: float # 吞吐 stopped: bool # 是否被用户中断 class ChatBackend(Protocol): """聊天后端接口。所有方法阻塞,调用方负责放 worker 线程。""" def load(self, on_status: Callable[[str], None]) -> None: """加载模型/权重;通过 on_status(msg) 上报进度。""" ... def generate( self, messages: list[dict], on_text: Callable[[str, int], bool], ) -> GenResult: """跑一轮生成。每步把「累计完整文本, 已生成 token 数」传给 on_text;返回 True 表示请求中断。""" ... @dataclass class FakeBackend: """测试用假后端:不加载模型,把预设回答按字符流式吐出。 delay > 0 时每字符间 sleep,用于 --demo 模式模拟真实吐字节奏;测试默认 0(不拖慢)。 """ reply: str = "你好，这是一个测试回答。" status_msgs: list[str] = field(default_factory=lambda: ["加载中(模拟)…"]) delay: float = 0.0 seen_messages: list[list[dict]] = field(default_factory=list) def load(self, on_status: Callable[[str], None]) -> None: for m in self.status_msgs: on_status(m) def generate(self, messages, on_text, tools=None, sampling=None) -> GenResult: import time self.seen_messages.append([dict(m) for m in messages]) acc = "" for ch in self.reply: acc += ch if self.delay: time.sleep(self.delay) # 用字符数近似 token 数(假后端无真实分词) if on_text(acc, len(acc)): return GenResult(acc, len(acc), 0.0, stopped=True) return GenResult(acc, len(self.reply), 0.0, stopped=False) def _common_prefix_len(a, b) -> int: """返回两个 token id 序列的最长公共前缀长度。""" n = min(len(a), len(b)) i = 0 while i < n and a[i] == b[i]: i += 1 return i def _reuse_prefix_len(cached_ids, new_ids) -> int: """可跨轮复用的前缀长度:仅当旧 cache 的 token 是新序列的**严格前缀**且新序列更长时, 返回该前缀长度(= len(cached_ids));否则返回 0 表示需全量重建。 只在严格前缀时复用,是为了永远「只延续、不回退」cache——Qwen3-Next 的线性注意力递归态 无法裁剪回任意历史位置;而 detokenize→retokenize 不一致、/reset、编辑历史等都会让公共前缀 短于旧长度,此时回退整段重建,绝不基于错位的 cache 续算。 """ if not cached_ids or len(new_ids) <= len(cached_ids): return 0 c = _common_prefix_len(cached_ids, new_ids) return c if c == len(cached_ids) else 0 class MLXBackend: """真实后端:封装 _build_engine + mtp_generate。MLX 相关 import 全部延迟到方法内。""" def __init__(self, args): self.args = args self._model = None self._tok = None self._drafter = None # 跨轮复用:持久化上一轮的 main_cache 及其对应的 token 序列(prompt + 已入 cache 的生成 token)。 self._main_cache = None self._cached_ids: list[int] = [] # 跨会话前缀快照:{head_key: (snaps, head_ids)},LRU。agentic 场景每个新会话 # 都重复同一段长 system+tools 前缀,快照后首轮 prefill 从 ~2600 tok 降到几百。 self._head_snaps: "OrderedDict[tuple, tuple]" = OrderedDict() # (system 文本, 工具名) → 固定头部 token 数,见 _fixed_prefix_len self._fixed_head_key = None self._fixed_head_len = 0 # 流式反分词器(单例,每轮 reset)。建实例要铺一张 15 万项的 id→token 表,不能每轮重建; # 生成由 server 的单飞锁串行,单例安全。 self._detok = None def load(self, on_status: Callable[[str], None]) -> None: from mlx_streaming.cli import _build_engine, _warmup self._model, self._tok, self._drafter = _build_engine( self.args, on_status=on_status ) # 预热:把首轮的 kernel 编译 + 专家池填充开销移到加载阶段,避免第一条消息莫名卡很久。 on_status("预热中(编译 kernel + 填专家池)…") _warmup(self._model, self._tok, self._drafter, self.args) # 采样参数:读模型自带 generation_config(Qwen3-Next 官方要求 do_sample # temp=0.7/top_p=0.8/top_k=20;贪心 argmax 会诱发幻觉/编造,见 2026-07 排查) import json import os self._gen_cfg = {"temperature": 0.7, "top_p": 0.8, "top_k": 20} try: with open(os.path.join(str(self.args.model), "generation_config.json")) as f: gc = json.load(f) if gc.get("do_sample", True): for k in ("temperature", "top_p", "top_k"): if k in gc: self._gen_cfg[k] = gc[k] except OSError: pass def _generate_sampled(self, ids, on_tokens, max_tokens, sampling, main_cache, cached_len): """采样生成(非投机):prefill 后逐 token 前向 + temp/top_p/top_k 采样。 返回 (produced, forwarded):forwarded = 已写入 cache 的生成 token 数, 供跨轮复用不变式(resident == len(ids)+forwarded)判定。 Qwen3-Next 的 generation_config 要求 do_sample;贪心 argmax 在 agentic 长 提示下会编造工具结果/幻觉工具名(实测)。代价:无 MTP 加速,decode 约 4-6 tok/s。 """ import mlx.core as mx from mlx_lm.sample_utils import make_sampler from mlx_streaming.mtp.generate import forward_with_hidden, prefill_chunked sampler = make_sampler( temp=float(sampling.get("temperature", 0.7)), top_p=float(sampling.get("top_p") or 0.0), top_k=int(sampling.get("top_k") or 0)) def _logprobs(logits): """末位 logits → float32 归一化 log-probs,这是 mlx_lm sampler 的入参契约。 必须归一化:apply_top_p 里是 `probs = mx.exp(x)` 再与 `1 - top_p` 比累积和, 喂未归一化的 logits 会让 exp() 变成天文数字、累积和瞬间越过阈值 → top_p 完全失效 (只砍掉概率 <1e-12 的尾巴)。必须转 float32:lm_head 输出是 bf16(8 位尾数), exp/cumsum/categorical 在 bf16 里做会把分布压得面目全非。 """ lg = logits[:, -1, :].astype(mx.float32) return lg - mx.logsumexp(lg, axis=-1, keepdims=True) ids_mx = ids if isinstance(ids, mx.array) else mx.array([ids]) logits, _ = prefill_chunked(self._model, ids_mx[:, cached_len:], main_cache) cur = _logprobs(logits) produced = [] forwarded = 0 for _ in range(max_tokens): nxt = int(sampler(cur)) produced.append(nxt) if on_tokens(produced[-1:]): break # 停止:当前 token 未写入 cache logits, _ = forward_with_hidden(self._model, mx.array([[nxt]]), main_cache) mx.eval(logits) forwarded += 1 cur = _logprobs(logits) return produced, forwarded # ---- 前缀快照磁盘持久化 ---- # 目录:models/prefix_snapshots/{key}.safetensors + {key}.json(sidecar:head ids/模型名)。 # 首个会话 prefill 一次落盘;之后(含引擎重启后)同前缀会话直接读盘恢复,消灭"首次"。 def _snap_dir(self) -> str: import os d = os.path.join(os.path.dirname(str(self.args.expert_dir)) or ".", "prefix_snapshots") os.makedirs(d, exist_ok=True) return d @staticmethod def _snap_key(ids, head: int) -> str: import hashlib import numpy as np return hashlib.sha1(np.array(ids[:head], dtype=np.int32).tobytes()).hexdigest()[:16] def _snap_save(self, key: str, snaps, ids, head: int) -> None: import json import os import mlx.core as mx from mlx.utils import tree_flatten from mlx_streaming import config as _cfg flat = {} metas = [] for ci, (st, meta) in enumerate(snaps): metas.append(list(meta) if isinstance(meta, (tuple, list)) else (meta or "")) for k, v in tree_flatten(st): flat[f"c{ci}.{k}"] = v path = os.path.join(self._snap_dir(), key) mx.save_safetensors(path + ".safetensors", flat) with open(path + ".json", "w") as f: # kv_quant 决定 cache 类型(量化/非量化),快照跨配置复用会 500 或错算,必须校验 json.dump({"head": head, "model": os.path.basename(str(self.args.model)), "kv_quant": _cfg.kv_quant(), "metas": metas, "ids": list(ids[:head])}, f) def _snap_load(self, key: str, ids, head: int): """磁盘命中返回 snaps(供 _cache_restore);缺失/模型不符/损坏返回 None。""" import json import os import mlx.core as mx from mlx.utils import tree_flatten, tree_unflatten path = os.path.join(self._snap_dir(), key) try: with open(path + ".json") as f: side = json.load(f) if side.get("head") != head or side.get("ids") != list(ids[:head]): return None if side.get("model") != os.path.basename(str(self.args.model)): return None # kv_quant 配置必须一致:量化 KV 的快照装进非量化 cache(或反之)会崩/错算 from mlx_streaming import config as _cfg if bool(side.get("kv_quant", True)) != bool(_cfg.kv_quant()): return None flat = mx.load(path + ".safetensors") metas = side.get("metas") or [] n_cache = max(int(k.split(".", 1)[0][1:]) for k in flat) + 1 snaps = [] for ci in range(n_cache): sub = {k.split(".", 1)[1]: v for k, v in flat.items() if k.startswith(f"c{ci}.")} meta = metas[ci] if ci < len(metas) else "" if isinstance(meta, list): meta = tuple(meta) snaps.append((tree_unflatten(list(sub.items())), meta)) mx.eval([v for st, _ in snaps for v in tree_flatten(st)]) return snaps except (OSError, ValueError, KeyError, json.JSONDecodeError): return None def _fixed_prefix_len(self, messages, tools) -> int: """算出 prompt 里「与用户问题无关」的固定头部长度(system + tools 渲染后的 token 数)。 做法是把同样的 system/tools 配一个哑 user 再渲染一次,取两次渲染的公共前缀—— 分叉点就是 user 内容的起始位置。这样得到的边界与提问内容无关,是快照能命中的最大头部: 再多一个 token 就把问题本身包进快照,换个问法即失效;少了则白白重算已知不变的部分 (实测 system+4 工具 = 3972 token,固定切 2048 会让每个新会话多 prefill 1900+ token)。 结果按 (system 文本, 工具名集合) 缓存:每轮只在配置变化时付一次 tokenize。 """ from mlx_streaming.cli import _encode_chat head_msgs = [m for m in messages if m.get("role") == "system"] if not head_msgs: return 0 key = ("".join(m.get("content") or "" for m in head_msgs), tuple(sorted((t.get("function") or {}).get("name", "") for t in (tools or [])))) if getattr(self, "_fixed_head_key", None) == key: return self._fixed_head_len probe = _encode_chat(self._tok, head_msgs + [{"role": "user", "content": "\x00probe"}], tools=tools) real = _encode_chat(self._tok, head_msgs + [{"role": "user", "content": "\x01other"}], tools=tools) n = 0 for a, b in zip(probe, real): if a != b: break n += 1 self._fixed_head_key, self._fixed_head_len = key, n return n def _snapshot_prefix(self, ids, fixed_head: int = 0): """跨会话前缀快照。内存命中/磁盘命中:恢复到新 cache 返回 (cache, head); 未命中:prefill 头部、深拷贝存内存 + 落盘后返回 (cache, head) (首个会话不增加前向量,只是把整段 prefill 拆两段);短 prompt 或关闭时返回 (None, 0)。 head 取 min(PREFIX_SNAPSHOT_HEAD 上限, system+tools 实际长度):快照必须停在 用户问题之前才能被后续会话命中,而停得越晚省下的 prefill 越多。 """ import mlx.core as mx from mlx_streaming import config from mlx_streaming.mtp.generate import prefill_chunked from mlx_streaming.mtp.kv_cache import _restore as _cache_restore from mlx_streaming.mtp.kv_cache import _snapshot as _cache_snapshot head = config.prefix_snapshot_head() if fixed_head > 0: head = min(head, fixed_head) head = min(head, len(ids) - 1) # 至少留 1 个 token 走前向,否则拿不到 logits # 门槛按 head 自身算,不按「prompt 比 head 长多少」:快照省下的就是 head 这段, # 与后面剩多少无关。头部 3972、问题只有 15 token 恰恰是收益最大的情形。 if head < 256: return None, 0 key = self._snap_key(ids, head) ent = self._head_snaps.get(key) if ent is not None: cache = self._model.make_cache() _cache_restore(cache, ent) self._head_snaps.move_to_end(key) return cache, head snaps = self._snap_load(key, ids, head) if snaps is not None: cache = self._model.make_cache() _cache_restore(cache, snaps) self._head_snaps[key] = snaps self._head_snaps.move_to_end(key) return cache, head cache = self._model.make_cache() prefill_chunked(self._model, mx.array([ids[:head]]), cache) snaps = _cache_snapshot(cache) self._head_snaps[key] = snaps self._head_snaps.move_to_end(key) while len(self._head_snaps) > config.prefix_snapshot_entries(): self._head_snaps.popitem(last=False) try: self._snap_save(key, snaps, ids, head) except Exception: # noqa: BLE001 落盘失败不影响主路径(内存快照仍生效) pass return cache, head def generate(self, messages, on_text, tools=None, sampling=None) -> GenResult: import time import mlx.core as mx from mlx_streaming.cli import _encode_chat, _eos_set, _truncate_eos from mlx_streaming.mtp.generate import mtp_generate tok = self._tok eos = _eos_set(tok) ids = _encode_chat(tok, messages, tools=tools) # 跨轮复用 KV/递归态:旧 cache 是本轮 prompt 的严格前缀时,只 prefill 新增后缀, # 不重算整段历史(prefill 从 ∝历史长度 降到 ∝新消息长度)。否则全量重建。 cached_len = (_reuse_prefix_len(self._cached_ids, ids) if self._main_cache is not None else 0) main_cache = self._main_cache if cached_len else None # 跨会话前缀快照:同一段长 system+tools 头部的新会话,跳过头部 prefill。 # 首个带该头部的会话在 prefill 头部后深拷贝一份状态存下(不产生额外前向, # 只是把一个整段 prefill 拆成两段),后续会话直接恢复。 if main_cache is None: main_cache, cached_len = self._snapshot_prefix( ids, fixed_head=self._fixed_prefix_len(messages, tools)) if main_cache is None: main_cache = self._model.make_cache() cached_len = 0 # 流式反分词:必须逐 token 喂 StreamingDetokenizer,不能每步 tok.decode(全部已产 token)。 # Qwen 是 byte-level BPE,很多汉字/符号跨 2 个 token,整体 decode 到半个字符时会按 # errors="replace" 吐出 U+FFFD;而下游按字符数算增量(delta = full[len(prev):]),下一步 # U+FFFD 被真字符替换后长度不变 → 增量为空、真字符永久丢失,客户端就留下一堆「�」。 # detokenizer 的 .text 只包含已完整的字符(未完成字节留在 _unflushed),恒为下一步的真前缀。 detok = self._detok if detok is None: detok = self._detok = tok.detokenizer detok.reset() produced_all: list[int] = [] # 已上报的正文 token(不含 EOS 及其后) stopped = {"v": False} hit_eos = {"v": False} def on_tokens(new_ids): for t in new_ids: if int(t) in eos: hit_eos["v"] = True break detok.add_token(int(t)) produced_all.append(int(t)) if hit_eos["v"]: detok.finalize() # 冲出尾部残留字节(正常收尾,不会有半字) if on_text(detok.text, len(produced_all)): # 用户按 Esc 请求中断 stopped["v"] = True return True # 命中 EOS:完整回答已生成,提前停止,避免引擎空跑到 max_tokens 让界面 # 长时间卡在「思考中」。EOS 属正常完成,不算中断。 return hit_eos["v"] t0 = time.perf_counter() if sampling is not None: # 采样模式(server 默认,质量优先):非投机逐 token 采样,无 MTP 加速。 produced, forwarded = self._generate_sampled( ids, on_tokens, self.args.max_tokens, sampling, main_cache, cached_len) stats = {"resident_tokens": len(ids) + forwarded} else: produced, stats = mtp_generate( self._model, self._drafter, tok, mx.array([ids]), self.args.max_tokens, K=self.args.k, ids_mode=True, profile=False, on_tokens=on_tokens, main_cache=main_cache, cached_len=cached_len, ) dt = time.perf_counter() - t0 # 持久化本轮 cache 供下轮复用。正常情况下 main_cache 恰好持有 `ids + produced[:-1]` # (produced[-1] 为 pending 未入 cache)。但末步多 token 跨 max_tokens 会 over-commit: # cache 领先于 produced,无法用已知 token 精确表述——此时禁用复用,下轮全量重建,绝不错算。 resident = stats.get("resident_tokens") expected = len(ids) + len(produced) - 1 if resident == expected: self._main_cache = main_cache self._cached_ids = list(ids) + list(produced[:-1]) else: self._main_cache = None self._cached_ids = [] out_ids = _truncate_eos(produced, eos) text = tok.decode(out_ids) tps = len(out_ids) / dt if dt > 0 else 0.0 return GenResult(text, len(out_ids), tps, stopped=stopped["v"])