437 lines
20 KiB
Python
437 lines
20 KiB
Python
"""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 被真字符替换后长度不变 → 增量为空、真字符永久丢失,客户端就留下一堆「<E5A086>」。
|
||
# 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"])
|