commit 4745f264b26705f8ce7f98072b47f37b5e18ea1a Author: fiser_jun Date: Tue Aug 4 14:34:00 2026 +0800 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..958e6c6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,40 @@ +.venv/ +__pycache__/ +*.pyc +.pytest_cache/ +.worktrees/ +*.gputrace/ + +# 模型权重本地放置（仅提交 models/README.md） +models/** +!models/README.md +*.safetensors +/experts/ + +# CMake 缓存 +native/ext/build/ + +# 日志与一次性产物 +logs/ +*.log +*.bak +benchmarks/reports/*.json +/dl_*.sh +/dl_*.log +/build_*.sh +/build_*.log +/extract_mtp.log +/dl_model.py + +# 根目录校验脚本（非运行时） +/verify_*.py +# 本地可留、不入库 +mlx_streaming/tests/ +mlx_streaming/tools/ +benchmarks/*.py +.DS_Store + +# 本地可留、不入库 +/性能优化与选型*.md +/8bit*.md +/演示*.md diff --git a/README.md b/README.md new file mode 100644 index 0000000..b3273c1 --- /dev/null +++ b/README.md @@ -0,0 +1,187 @@ +# Sparkle + +本地 **Qwen3-Next-80B-A3B Instruct 8bit** 流式推理引擎，对接 ccSparkle Agent（模型源：**Sparkle**）。 + +在约 64GB 统一内存的 Apple Silicon 上，把约 77GB 专家权重放在 SSD，运行时按需装入 + 预取，从而跑通 80B 8bit（oMLX / Ollama 全量加载做不到的档位）。动机见 [`models/README.md`](./models/README.md)：用 80B@64G 验证流式路径，对应下游 **30B 8bit A3B** 进 32G/16G、统一内存省 50%+。 + +| 项 | 值 | +|----|----| +| Python | **3.14.***（与 `native_moe_ext.cpython-314-darwin.so` 对齐；见 `pyproject.toml`） | +| 已验证机器 | MacBook Pro M2 Max / 64GB | +| 安装目录 | 本仓库根目录（下文记为 `$SPARKLE_HOME`，即 `clone` 后的路径） | +| 引擎端口 | `8317`（经中台 `8088` 启停与代理） | +| 模型 / 专家 / MTP | 见 [`models/README.md`](./models/README.md)（不入库，约 160GB） | +| 性能基线 | 稳态约 9–10 tok/s（贪心+MTP）；采样约 4–6 tok/s；常驻约 23–33GB 量级 | + +依赖安装走清华 PyPI 镜像（`pyproject.toml` 的 `tool.uv.index`）。 + +--- + +## 1. 系统组成 + +``` +浏览器前端(8090) ← 日常对话界面 + │ 顶部选 Sparkle 源 + ▼ +中台代理(8088) ← 转发请求，规避跨域 + ▼ +Sparkle 引擎(8317) ← 80B 8bit 流式推理 + ├─ /v1/chat/completions OpenAI 兼容接口 + └─ /admin 调参页（槽位 / 投机宽度 / 状态） +``` + +专家权重留在 SSD，运行时按需读取 + 预测式预取；首次启动和切换话题时偶有停顿（读盘）。 + +**AUTOPIN**（默认开启）：记录专家路由热度（`models/experts_8bit_g64/pool_usage.json`），启动时把最热一批 pin 进常驻池（默认占槽位 50%，定死不驱逐）。冷启动命中率可从约 72% 升到 **98%+**。`AUTOPIN=0` 关闭；`AUTOPIN_BUDGET_FRAC` 调比例（0–1）。 + +--- + +## 2. 启动 + +**推荐**：打开对话界面 → 顶部选 **Sparkle** → 中台自动加载引擎（约 90 秒，状态变绿后再对话）。 + +1. 中台（桌面应用常会自动拉起；路径按你本机 Agent 工程根目录，下文记 `$AGENT_HOME`）： + ```bash + cd "$AGENT_HOME" && node scripts/zhongtai-server.js + ``` +2. 前端静态服务： + ```bash + cd "$AGENT_HOME/<前端静态目录>" && python3 -m http.server 8090 + ``` +3. 浏览器打开对话页，顶部点 **Sparkle**——未运行时会自动启动（也可在 `配` → Sparkle tab「引擎管理」点**加载**）。 + +就绪：`配` → Sparkle 状态点变绿，或 `curl -s http://127.0.0.1:8088/sparkle-engine/status` 见 `"loaded":true`。 + +**手动启动（调试）**： + +```bash +cd "$SPARKLE_HOME" # 本仓库根目录 +.venv/bin/python -m mlx_streaming.server --port 8317 \ + --model models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit \ + --expert-dir models/experts_8bit_g64 \ + --qn-config models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit/config.json \ + --mtp-out models/qn_mtp_weights.safetensors \ + --expert-slots 96 +``` + +- Admin：`http://127.0.0.1:8317/admin` +- 健康：`curl -s http://127.0.0.1:8317/api/stats`（`unplaced_experts` 应为 `0`） + +中台默认 `SPARKLE_ENGINE_DIR` → 本目录。改目录名后若「卸载」失灵，先杀 `8317` 再重启中台。 + +> **工具白名单**：默认只放画像 4 工具（`profile_list_employees` / `profile_get_portrait_overview` / `profile_get_person_report` / `profile_get_team_report`）。中台动态工具集会改 system 头部哈希，导致前缀快照 miss；钉死工具集后快照才稳定。排班 `metro_*` 等会被裁掉。临时放开：`TOOLS_ALLOW=`；换名单：`TOOLS_ALLOW=a,b,c`。裁剪结果见 `logs/server.log` 的 `[TOOLS_ALLOW]`。 + +> **不要自己设 `PREFILL_CHUNK`**。安全上限由槽位决定；引擎启动时按槽位自算（96 槽 → 4，64 槽 → 3），超限会被钳制并打日志。 + +--- + +## 3. 日常对话 + +1. 打开对话页（`http://localhost:8090`），顶部选 **Sparkle** +2. 首次：`配` → Sparkle tab：Base URL `http://127.0.0.1:8317`，API Key 任意非空（如 `sparkle`），刷新选模型 `Qwen3-Next-80B-A3B-Instruct-MLX-8bit`，保存 +3. 正常对话；工具调用链路与 oMLX 源相同 + +--- + +## 4. 调参与引擎管理 + +`配` → **本地 Sparkle** tab「引擎管理」（与 `http://127.0.0.1:8317/admin` 等价）： + +状态点：灰（未启动）→ 黄（加载中，约 90 秒）→ 绿（就绪）。`加载` / `卸载` 启停；就绪时显示 tok/s 与 RSS。 + +| 参数 | 范围 | 说明 | 生效 | +|------|------|------|------| +| expert-slots | 32–160 | 常驻专家数，越大命中率越高、越占内存（每槽 ~160MB）。64GB + AUTOPIN：64 槽约 9 tok/s；**96 槽约 10.3 tok/s（推荐）**；≥128 近红线；**160 会卡死，勿用** | 应用后进程级重启（约 90 秒） | +| K（投机宽度） | 1–5 | MTP 草稿 token 数，默认 3；只影响速度 | 即时 | +| max_tokens | — | 单次回答最大长度 | 即时 | + +日常 64 够用；长文档可试 96–128；系统开始 swap 就降回 64。 + +--- + +## 5. 故障排查 + +| 现象 | 原因 | 处理 | +|------|------|------| +| 503 / engine reloading | 加载或重建中 | 等 60–90 秒后再试 | +| `EADDRINUSE ... 8317` | 旧 server 未关 | `lsof -tiTCP:8317 -sTCP:LISTEN \| xargs kill` | +| 中台 `EADDRINUSE ... 8088` | 已有健康中台 | 直接用；要重启先 kill | +| 改名后「卸载」无效 | 旧进程挂在旧路径 | 杀 8317，重启中台再加载 | +| 明显变慢 | swap 或槽位太低 | 看内存压力；调槽；关其他大应用 | +| 首轮特别慢 | 冷启动 | 正常，第二轮起恢复 | +| 长文中途停顿 | 读盘装冷门专家 | 正常；调高槽位可缓解 | +| 刷新拉不到模型 | server 未起或端口错 | `curl http://127.0.0.1:8317/v1/models -H "Authorization: Bearer x"` | + +--- + +## 6. 常见问题 + +**Q: 为什么 oMLX / Ollama 里看不到这个 8bit 模型？** +A: 权重远超 64GB 整模上限，只有 Sparkle 流式能跑。oMLX 继续跑中小模型即可。 + +**Q: 投机宽度 K 是什么？** +A: MTP 每步草稿 token 数；主模型批量验证，猜中白赚。只影响速度，不影响内容。 + +**Q: 回答乱编 / 不调工具 / 幻觉？** +A: Qwen3-Next 官方要求采样（temp 0.7 / top_p 0.8 / top_k 20）。贪心+MTP 在该模型上易幻觉。默认走采样（约 4–6 tok/s）；显式 `temperature: 0` 才走贪心+MTP（约 10 tok/s）。Sparkle 客户端对真实对话会把 `0` 改回采样。 + +**Q: 模块分析比 oMLX 7B 还空？** +A: 多为分析轮传了 `temperature: 0` 进贪心导致指令遵循崩坏，不是花名册数据错。另查 `unplaced_experts`（非 0 = 错专家）与是否需清空 `prefix_snapshots`。 + +**Q: 回答里夹「�」乱码？** +A: 已修：byte-level BPE 流式反分词用 `StreamingDetokenizer`，避免半字符 U+FFFD 丢字。 + +**Q: 8bit 比 7B 还笨？** +A: 曾因 prefill 超槽静默映射到 0 号专家。现自动算 `PREFILL_CHUNK`、超容量走正确慢路径，并用 `unplaced_experts` 告警（健康值恒为 0）。非 0 时回答不可信；若错算过，删 `models/prefix_snapshots/` 再预热。 + +**Q: 和云端略有不同？** +A: 量化 + KV 压缩的正常差异；语义质量不受影响。 + +**Q: 首轮为什么要等？** +A: 新对话要 prefill 系统提示+工具定义。有前缀快照后：重启后首会话约百秒级，之后新会话约数十秒，同会话追问约数秒。`PREFIX_SNAPSHOT_HEAD` 调长度（默认 2048，0 关闭）。 + +**Q: 速度还能再快吗？** +A: 96 槽 + AUTOPIN + prefill 分块约 9–10 tok/s，接近本机 8bit 上限。160 槽会卡死，勿试。 + +**Q: 能跑多长上下文？** +A: KV 量化默认开；对话侧另有约 24k token 压缩阈值（中台配置）。 + +**Q: 想换模型或其它量化档？** +A: 本引擎面向 Qwen3-Next-80B 流式 8bit 部署；换档需按当前模型重新准备专家 blob 与 MTP。 + +--- + +## 7. 目录要点 + +``` +sparkle/ +├── mlx_streaming/ # 引擎 Python 包（含已编译 native_moe_ext*.so） +├── native/ext/ # C++/Metal 扩展源码（重编 .so 用） +├── models/ # 权重与专家 blob（约 160GB，勿删） +│ ├── Qwen3-Next-80B-A3B-Instruct-MLX-8bit/ +│ ├── experts_8bit_g64/ +│ └── qn_mtp_weights.safetensors +├── .venv/ +├── pyproject.toml / uv.lock +└── benchmarks/reports/ # 性能基线报告 +``` + +`models/` 三者缺一不可。不要把模型目录移进 oMLX 等其它运行时的模型扫描目录（会列出但整模加载不了）。 +换 Python 小版本时需 `make -C native/ext native_moe_ext` 重编 `.so`。`native/ext/build/` 已忽略。 + +--- + +## 8. 质量与速度（摘要） + +- **默认应对话走采样解码**（约 temperature 0.7 / top_p 0.8 / top_k 20）。`temperature=0` 会进贪心+MTP，本机实测指令遵循差于同表 oMLX 7B。 +- 稳态 decode：采样约 **4–6 tok/s**；贪心+MTP 约 **9–10 tok/s**（速度更高，但质量风险大，生产默认勿用）。 +- 专家池命中率健康时约 **94%+**；`unplaced_experts ≠ 0` 表示曾用错专家权重，回答不可信。 + +--- + +## 9. 相关文档 + +| 文档 | 内容 | +|------|------| +| [models/README.md](./models/README.md) | 为何做流式、NAS 下载与专家/MTP 准备 | +| [benchmarks/reports/8bit-g64-baseline-2026-07-25.md](./benchmarks/reports/8bit-g64-baseline-2026-07-25.md) | 8bit 性能基线 | diff --git a/benchmarks/reports/8bit-g64-baseline-2026-07-25.md b/benchmarks/reports/8bit-g64-baseline-2026-07-25.md new file mode 100644 index 0000000..7b8d775 --- /dev/null +++ b/benchmarks/reports/8bit-g64-baseline-2026-07-25.md @@ -0,0 +1,61 @@ +# Qwen3-Next-80B-A3B 8bit(g64)基线 — M2 Max 64GB + +日期:2026-07-25 · 机器:MacBook Pro M2 Max / 64GB / 内置 SSD +模型:`models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit`(**bits=8, group_size=64**,79GB) +数据:`models/experts_8bit_g64/blobs`(48 层,77GB,v1 affine,stride 3,342,336B,page-aligned)+ `models/qn_mtp_weights.safetensors`(3.3GB) + +## 性能 + +| 场景 | tok/s | accept_len | 备注 | +|---|---|---|---| +| 冷启动首轮(32 tok) | 5.0 | 1.889 | 含 kernel 编译 + 池冷填充 | +| 稳态(128 tok,64 槽) | **9.4** | 2.246 | 符合 8-12 tok/s 预期 | +| server 短请求(33 tok,含 prefill) | 6.6 | — | OpenAI 链路实测 | + +- 常驻内存:RSS ≈ **23.5 GB**(`--expert-slots 64`,池 ~10GB + dense + 缓冲;64GB 机器余量充足,可上探 96-128 槽) +- 加载时间:约 60-90s(含预热) + +## 正确性(三层证据链) + +1. **kernel 单元等价**(`verify_8bit_equiv.py`):L0/L15/L47 blob fused kernel vs Python 参考 **diff = 0** +2. **端到端对拍**(`verify_e2e_spec.py`):贪心 vs 投机 token 17 确定性分叉 → 触发对照 +3. **KV 量化对照**(`verify_teacher_forced.py`,seq=1 vs seq=2 生产原语): + - `KV_QUANT=0`:**24/24 token 全一致**,cosine 0.99935 → 8bit 权重路径数值干净 + - `KV_QUANT=1`(生产默认):23/24(95.8%),cosine 0.99870,翻转处 margin=0 近 tie → 达质量门(≥95%/≥0.99) + - 结论:分叉全部归因于 KV 量化(KV 量化既有有损特性),与权重量化位宽无关 + +## 已知事项 + +- `VirtualPool.acquire_host` latent bug:单次前向单层唯一专家数超池容量时 fetch 回退会 AttributeError;生产 prefill 分块 chunk=2 永不触发。勿用整段 seq=N 一次性前向做诊断。 +- server `/admin` 调参页池内存估算已按 g64 修正为 ~160MB/槽。 +- server 模型 id 取目录名:`Qwen3-Next-80B-A3B-Instruct-MLX-8bit`。 + +## 复现命令 + +```bash +# CLI +.venv/bin/sparkle --model models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit \ + --expert-dir models/experts_8bit_g64 \ + --qn-config models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit/config.json \ + --mtp-out models/qn_mtp_weights.safetensors --expert-slots 64 --stats + +# OpenAI server(前端 Sparkle 源对接,端口 8317) +.venv/bin/python -m mlx_streaming.server --port 8317 \ + --model models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit \ + --expert-dir models/experts_8bit_g64 \ + --qn-config models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit/config.json \ + --mtp-out models/qn_mtp_weights.safetensors --expert-slots 64 +``` + +## 追记(2026-07-25 晚):AUTOPIN + 调优后 + +| 配置 | decode tok/s | 备注 | +|---|---|---| +| 64 槽 + AUTOPIN | ~8.0-9.4 | 稳定,RSS ~24GB | +| **96 槽 + AUTOPIN(生产)** | **9.3-10.3** | 甜点位,RSS ~32GB;CLI 128tok: 9.3(MTP4)/9.7(MTP8) | +| 160 槽 + AUTOPIN | 不可用 | RSS ~34GB 系统耗尽卡死(AUTOPIN eager 预热多占 ~5GB) | + +- AUTOPIN(热度持久化+启动 pin 定死):命中率 72.2%→98.6%,pin 3439 专家/6.5s +- PREFILL_CHUNK=8:676 tok prompt 2m49s→33s +- MTP_BITS 4→8 A/B:accept_len 2.246→2.207,无效,保持 4 +- 槽位调整改进程级重启(进程内重建在新 C++ pin 状态下会把引擎拖到 0.24 tok/s) diff --git a/mlx_streaming/.gitignore b/mlx_streaming/.gitignore new file mode 100644 index 0000000..77ac754 --- /dev/null +++ b/mlx_streaming/.gitignore @@ -0,0 +1,3 @@ +.venv/ +__pycache__/ +*.pyc diff --git a/mlx_streaming/__init__.py b/mlx_streaming/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mlx_streaming/cli.py b/mlx_streaming/cli.py new file mode 100644 index 0000000..179ac97 --- /dev/null +++ b/mlx_streaming/cli.py @@ -0,0 +1,329 @@ +"""sparkle 命令行:面向用户的交互式多轮对话(MTP 自投机快路径)。 + +用法示例: + sparkle # 直接进入交互式对话(默认子命令 chat) + sparkle chat # 同上 + sparkle -k 4 -n 800 --stats # 调宽投机、加长生成、每轮打印吞吐 + sparkle --system "你是一个简洁的助手" + sparkle --model models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit --expert-slots 32 + +只做「生成」一件事:走 MTP 自投机 + 零拷贝双源侧区快路径。关键参数做成命令行 flag, +其余调优项仍从环境变量读取(见 mlx_streaming/config.py)。 + +交互期间可用命令: + /exit 或 /quit 退出 + /reset 清空对话历史(保留 system) + /help 打印帮助 +""" +import argparse +import sys +import time + +from mlx_streaming import config + +# MTP 快路径环境变量兜底配方(benchmark 验证过的最优组合)。 +# 用 setdefault 兜底:用户显式导出的环境变量优先级更高,不会被覆盖。 +# MTP_ADAPTIVE_DEPTH:置信度门控动态深度。逐位累计置信度跌破 tau 即停,低置信步抽浅省专家加载 +# (本系统 IO 瓶颈)。消融(reports/adaptive-depth-2026-07-05)证 τ=0.3、depth_max=3 纯向下收缩 +# +5~6% tok/s 且 bit-lossless、零额外显存,稳定优于最小树 pos0 救回 → 设为用户主路径默认。 +# depth_max=3 与基础 K 一致,在生产 EXPERT_SLOTS=32 下 seq·top_k 不溢出 cap(扩到 4 须 slots>=40)。 +# 注:动态深度与 TREE_TOP2 互斥(adaptive 仅在非 tree 的 plain 路径生效),故此处不开 TREE_TOP2。 +# KV_QUANT:IsoQuant K4/V3 + SO(4) 块旋转,仅作用于 12 个全注意力层(线性层递归态不动)。 +# 极致压缩长上下文 KV:128k 3.0→~0.68 GiB;短会话收益小但无害。K4/V3/旋转均取默认值, +# 开 KV_QUANT=1 即整套生效。质量验收:token 一致率≥95% + logits cosine≥0.99。 +_FASTPATH_ENV = { + "STREAM_BLOB_LOADER": "1", + "NATIVE_FUSED_PREFETCH": "1", + "ZEROCOPY_DUAL_SOURCE": "1", + "SIDEREGION_LFU": "1", + "KV_QUANT": "1", + "MTP_ADAPTIVE_DEPTH": "1", + "MTP_CONF_TAU": "0.3", + "MTP_DEPTH_MAX": "3", +} + + +def _build_engine(args, on_status=None): + """按 MTP 快路径装配 model / tokenizer / drafter。 + + 注意:model_builder 在 import 时就读 MODEL/EXPERT_DIR 等环境变量,所以必须 + 先把命令行参数写进 os.environ,再 import build_streaming_model。 + + on_status:可选进度回调;为 None 时进度打到 stderr(保持旧行为)。 + """ + def _emit(msg): + if on_status is not None: + on_status(msg) + else: + print(msg, file=sys.stderr, flush=True) + + import os + + os.environ["MODEL"] = args.model + os.environ["QN_CONFIG"] = args.qn_config + os.environ["MTP_OUT"] = args.mtp_out + os.environ["EXPERT_DIR"] = args.expert_dir + os.environ["EXPERT_SLOTS"] = str(args.expert_slots) + spec = args.spec_slots if args.spec_slots is not None else args.expert_slots + os.environ["POOL_SPEC_SLOTS"] = str(spec) + for k, v in _FASTPATH_ENV.items(): + os.environ.setdefault(k, v) + + import json + + import mlx.core as mx # noqa: F401 确保 MLX 已就绪 + from mlx_lm.models.qwen3_next import ModelArgs + + from mlx_streaming.core.mem import setup_memory_hygiene + from mlx_streaming.model_builder import build_streaming_model + from mlx_streaming.mtp.drafter import MTPDrafter + from mlx_streaming.mtp.qwen3_next_mtp import load_mtp + + # 长会话内存防御:封顶 MLX 可回收缓冲(默认 1GB),防长对话里缓冲缓存膨胀把常驻推过墙 / + # 触发 macOS 压缩器抖动。TUI 是典型长会话场景,故在用户主路径启动时就设上。 + _applied = setup_memory_hygiene(cache_gb=config.mlx_cache_limit_gb(), + wired_gb=config.mlx_wired_limit_gb()) + if _applied: + _emit(f"内存防御: {_applied}") + + _emit("正在加载主模型 + 专家(流式)...") + model, tok, _store = build_streaming_model() + _emit("正在加载 MTP drafter...") + with open(args.qn_config) as f: + margs = ModelArgs.from_dict(json.load(f)) + mtp = load_mtp(margs, args.mtp_out, quantize=True, bits=config.mtp_bits()) + mtp.embed_tokens = model.model.embed_tokens # 共享主模型 embedding + drafter = MTPDrafter(mtp, model.lm_head) + return model, tok, drafter + + +def _warmup(model, tok, drafter, args): + """跑一次生成做预热:首轮的明显卡顿主要来自现编译 Metal kernel + 填 MoE 专家 resident 池, + 提前把这部分一次性开销移到加载阶段。 + + 覆盖增强:用一段**较长、token id 跨大跨度词表分散**的合成 prompt——MoE 路由依赖 token 内容, + 分散的 id 会命中更多专家、更充分地预填专家池;较长 prompt 又能走通多块分块 prefill。 + 合成 id 直接走 ids_mode(不依赖 tokenizer),分块 prefill(chunk=2)保证长 prompt 也不抬高显存峰值。 + 预热失败不致命(直接吞掉异常),不影响后续真实生成。 + """ + import mlx.core as mx + + from mlx_streaming.mtp.generate import mtp_generate + + try: + vocab = int(model.model.embed_tokens.weight.shape[0]) + n = min(64, vocab) # 预热 prompt 长度(兼顾覆盖与耗时) + step = max(1, vocab // n) # 在词表内均匀取样,最大化专家覆盖 + ids = [(1 + i * step) % vocab for i in range(n)] + mtp_generate(model, drafter, tok, mx.array([ids]), 8, + K=args.k, ids_mode=True) + except Exception: # noqa: BLE001 预热仅为压首轮延迟,失败不应中断启动 + pass + + +def _encode_chat(tok, messages, tools=None): + """把多轮对话按聊天模板编码成 token id 列表。tools(OpenAI 格式)交给模板渲染工具定义。""" + tmpl = getattr(tok, "chat_template", None) + if tmpl: + kwargs = {"add_generation_prompt": True} + if tools: + kwargs["tools"] = tools + out = tok.apply_chat_template(messages, **kwargs) + # 返回形态随 transformers 版本而变,三种都要兜住: + # transformers 4.x → list[int] + # transformers 5.x(默认) → BatchEncoding,即 {"input_ids": [[...]], ...} + # 部分实现 → str + # 只写 list(out) 会在 5.x 下取到字典的**键名** ["input_ids", "attention_mask"], + # 于是 prompt 退化成两个字符串:主生成路径(backend.py 的 ids)喂进去全是乱码, + # 且 _fixed_prefix_len 恒为 2 → head < 256 → 前缀快照永久失效。 + # 单元测试全部 monkeypatch 掉了本函数,捕不到这个回归,故在此显式判型。 + if isinstance(out, str): + return tok.encode(out) + ids = out if isinstance(out, (list, tuple)) else out["input_ids"] + if len(ids) > 0 and isinstance(ids[0], (list, tuple)): + ids = ids[0] # 批次维度 + return [int(t) for t in ids] + # 无聊天模板:朴素拼接兜底。 + text = "" + for m in messages: + text += f"{m['role']}: {m['content']}\n" + text += "assistant: " + return tok.encode(text) + + +def _eos_set(tok): + """收集所有可能的结束符 token id。""" + eos = set() + ids = getattr(tok, "eos_token_ids", None) + if ids: + eos |= set(ids) + one = getattr(tok, "eos_token_id", None) + if one is not None: + eos.add(one) + return eos + + +def _truncate_eos(produced, eos): + """遇到第一个结束符即截断(不含结束符本身)。""" + for i, t in enumerate(produced): + if t in eos: + return produced[:i] + return produced + + +_HELP = """可用命令: + /exit, /quit 退出 + /reset 清空对话历史(保留 system) + /help 显示本帮助 +直接输入文本即可对话。""" + + +def cmd_chat(args): + """默认启动全屏 TUI;--plain 走纯文本 REPL;--demo 用假后端免模型预览界面。""" + if getattr(args, "plain", False): + return _chat_repl(args) + from mlx_streaming.tui import run_tui + if getattr(args, "demo", False): + # 免模型预览:秒开 TUI,假流式回答,用于验证界面/占位符/状态栏 + from mlx_streaming.tui.backend import FakeBackend + demo = FakeBackend( + reply="这是 --demo 测试回答:界面、逐字流式、状态栏(token 数 / tok·s)" + "均为模拟,不加载模型。按 Esc 可中断,/help 看命令。", + delay=0.03) + return run_tui(demo, args) + from mlx_streaming.tui.backend import MLXBackend + return run_tui(MLXBackend(args), args) + + +def _chat_repl(args): + model, tok, drafter = _build_engine(args) + + import mlx.core as mx + + from mlx_streaming.mtp.generate import mtp_generate + from mlx_streaming.tui.backend import _reuse_prefix_len + + eos = _eos_set(tok) + base_messages = [] + if args.system: + base_messages.append({"role": "system", "content": args.system}) + messages = list(base_messages) + # 跨轮复用:持久化上一轮 main_cache 及其对应的 token 序列(与 MLXBackend 同机制)。 + main_cache = None + cached_ids: list[int] = [] + + print("\n正在预热(编译 kernel + 填专家池)…", file=sys.stderr, flush=True) + _warmup(model, tok, drafter, args) + print("模型已就绪。输入 /help 查看命令,/exit 退出。", file=sys.stderr, flush=True) + + while True: + try: + user = input("\n你 > ").strip() + except (EOFError, KeyboardInterrupt): + print() + break + if not user: + continue + if user in ("/exit", "/quit"): + break + if user == "/reset": + messages = list(base_messages) + main_cache, cached_ids = None, [] # 清历史同时弃用旧 cache + print("对话历史已清空。", file=sys.stderr) + continue + if user == "/help": + print(_HELP, file=sys.stderr) + continue + + messages.append({"role": "user", "content": user}) + ids = _encode_chat(tok, messages) + + # 旧 cache 是本轮 prompt 的严格前缀时只 prefill 新增后缀,否则全量重建。 + cached_len = (_reuse_prefix_len(cached_ids, ids) + if main_cache is not None else 0) + cur_cache = main_cache if cached_len else model.make_cache() + + # EOS 提前停止:否则会空跑到 max_tokens,且 produced 会含 EOS 后垃圾 token, + # 使 cached_ids 与下轮编码前缀断裂、复用永不触发。 + produced_all: list[int] = [] + + def _on_tokens(new_ids): + produced_all.extend(new_ids) + truncated = _truncate_eos(produced_all, eos) + return len(truncated) < len(produced_all) + + t0 = time.perf_counter() + produced, stats = mtp_generate( + model, drafter, tok, mx.array([ids]), + args.max_tokens, K=args.k, ids_mode=True, profile=args.stats, + on_tokens=_on_tokens, main_cache=cur_cache, cached_len=cached_len) + dt = time.perf_counter() - t0 + + # offset 对账:仅无 over-commit 时记录 cache 供下轮复用(见 MLXBackend.generate)。 + if stats.get("resident_tokens") == len(ids) + len(produced) - 1: + main_cache = cur_cache + cached_ids = list(ids) + list(produced[:-1]) + else: + main_cache, cached_ids = None, [] + + out_ids = _truncate_eos(produced, eos) + text = tok.decode(out_ids) + print(f"\n助手 > {text}") + messages.append({"role": "assistant", "content": text}) + + if args.stats: + tps = len(out_ids) / dt if dt > 0 else 0.0 + print(f"[{len(out_ids)} tok, {tps:.1f} tok/s, " + f"accept_len={stats.get('avg_accept_len')}]", file=sys.stderr) + + print("再见。", file=sys.stderr) + return 0 + + +def _add_chat_args(p): + p.add_argument("--model", default=config.model_path(), help="主模型路径(MLX 量化)") + p.add_argument("--expert-dir", default=config.expert_dir(), + help="拆分后的 per-expert 目录") + p.add_argument("--mtp-out", default=config.mtp_out(), help="MTP 权重文件") + p.add_argument("--qn-config", default=config.qn_config(), + help="Qwen3-Next 配置 JSON") + p.add_argument("-k", "--k", type=int, default=3, help="MTP 投机宽度(默认 3)") + p.add_argument("-n", "--max-tokens", type=int, default=4096, + help="每轮最多生成的新 token 数(默认 4096)") + p.add_argument("--expert-slots", type=int, default=32, + help="常驻专家池容量(默认 32,同时作为侧区行数默认)") + p.add_argument("--spec-slots", type=int, default=None, + help="侧区行数 POOL_SPEC_SLOTS(默认跟随 --expert-slots)") + p.add_argument("--system", default=None, help="可选 system 提示词") + p.add_argument("--stats", action="store_true", + help="每轮结束在 stderr 打印 token 数 / tok·s / 接受长度") + p.add_argument("--plain", action="store_true", + help="用纯文本 REPL,不启动全屏 TUI(终端不兼容/调试时用)") + p.add_argument("--demo", action="store_true", + help="免模型预览:用假后端秒开 TUI,验证界面/流式/状态栏") + p.set_defaults(func=cmd_chat) + + +def _build_parser(): + parser = argparse.ArgumentParser( + prog="sparkle", + description="sparkle:Apple Silicon 上的流式 MoE + Qwen3-Next MTP 自投机推理") + sub = parser.add_subparsers(dest="cmd") + chat = sub.add_parser("chat", help="进入交互式多轮对话(MTP 自投机快路径)") + _add_chat_args(chat) + return parser + + +def main(argv=None): + argv = list(sys.argv[1:] if argv is None else argv) + subcmds = {"chat"} + # 让 chat 成为默认子命令:不带子命令(或首参是 flag)时自动补上 chat; + # 但保留顶层 -h/--help 直接显示总帮助。 + if not argv or (argv[0] not in subcmds and argv[0] not in ("-h", "--help")): + argv = ["chat"] + argv + args = _build_parser().parse_args(argv) + return args.func(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/mlx_streaming/config.py b/mlx_streaming/config.py new file mode 100644 index 0000000..4d46ae0 --- /dev/null +++ b/mlx_streaming/config.py @@ -0,0 +1,298 @@ +"""集中配置：所有运行时开关从环境变量读取（env 名与默认值与历史一致，便于现有跑法）。 + +设计目的：把原先散落在 core/ 各文件的 `os.environ.get` 全部收口到这里，单点可查、 +避免默认值漂移。读取均为运行时（每次调用读 env），与原行为一致——测试/probe 用 +monkeypatch env 仍生效，热路径每 token 读 env 的成本与原先相同。 + +对「同一 env 名在不同调用方有不同默认」的项（如 CROSS_LAYER_PREFETCH_AHEAD：hook=0、 +native 预取=1），accessor 暴露 `default` 参数，由调用方传入，绝不擅自统一。 +""" +import os + + +def _b(name: str, default: str = "0") -> bool: + return os.environ.get(name, default) == "1" + + +def _i(name: str, default) -> int: + return int(os.environ.get(name, str(default))) + + +def _f(name: str, default) -> float: + return float(os.environ.get(name, str(default))) + + +def _s(name: str, default: str = "") -> str: + return os.environ.get(name, default) + + +# ============================ 模型 / 目录 ============================ +# 默认放持久目录 models/（不要放 /tmp，会被系统清空）。运行从仓库根目录起，相对路径即可。 +def model_path() -> str: return _s("MODEL", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit") +def qn_config() -> str: return _s("QN_CONFIG", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit/config.json") +def mtp_out() -> str: return _s("MTP_OUT", "models/qn_mtp_weights.safetensors") +def expert_dir(default: str = "models/experts_8bit_g64") -> str: return _s("EXPERT_DIR", default) +def expert_slots() -> int: return _i("EXPERT_SLOTS", 64) +# 长期运行内存防御:封顶 MLX 可回收缓冲(默认 1GB),防长会话缓存膨胀; +# 双源侧区池的专家 buffer 走 C++ owned pool、不经 MLX 缓冲缓存,故 1GB 缓冲复用额度已够, +# 再大只是白占常驻。wired limit 默认 0=关(opt-in),设 >0 则 wire 该 GB 数的 GPU 缓冲防 macOS +# 压缩器,须 < 系统建议工作集(本机 26.8GB)。 +def mlx_cache_limit_gb() -> float: return _f("MLX_CACHE_LIMIT_GB", 1.0) +def mlx_wired_limit_gb() -> float: return _f("MLX_WIRED_LIMIT_GB", 0.0) +def expert_pool_profile() -> str: return _s("EXPERT_POOL_PROFILE", "") +def hidden_variant() -> str: return _s("HIDDEN_VARIANT", "pre_final_norm") +def blob_dir() -> str: return _s("BLOB_DIR", "") +def compute_buffer_dir() -> str: return _s("COMPUTE_BUFFER_DIR", "") + + +# ============================ MoE 热路径 ============================ +def resident_pool_enabled() -> bool: return _b("RESIDENT_POOL", "1") +def gpu_remap_enabled() -> bool: return _b("GPU_REMAP", "1") +# 实验：让 verify(seq>1)也走 GPU 侧 slot 重映射(acquire_gpu),消掉 host 路径每层 .tolist() 栅栏。 +def verify_gpu_remap() -> bool: return _b("VERIFY_GPU_REMAP", "0") +# A2：GPU 重映射路径下 promote 是否用 GPU membership 现算 used 过滤假阳性（默认开；回退设 0）。 +def gpu_remap_promote_filter() -> bool: return _b("GPU_REMAP_PROMOTE_FILTER", "1") +def moe_topk_override() -> "str | None": return os.environ.get("MOE_TOPK_OVERRIDE") +def eager_expert_load() -> bool: return _b("EAGER_EXPERT_LOAD", "0") + + +# ============================ 缓存 / 驱逐 ============================ +def evict_policy() -> str: return _s("EVICT_POLICY", "lfu") +# 专家读盘是否绕过 OS page cache(F_NOCACHE):默认开,保证基准每次都是真实 NVMe 读、 +# 结果可复现(不被页缓存冷热污染);EXPERT_NOCACHE=0 恢复 mmap/page-cache(重复跑更快但飘)。 +def expert_nocache() -> bool: return _b("EXPERT_NOCACHE", "1") +def lfu_decay_interval() -> int: return _i("LFU_DECAY_INTERVAL", 0) +def expert_bundle() -> bool: return _b("EXPERT_BUNDLE", "0") +def expert_bundle_dir(default: str) -> str: return _s("EXPERT_BUNDLE_DIR", default) +def expert_bundle_cache() -> int: return _i("EXPERT_BUNDLE_CACHE", 4) +def expert_stack_cache() -> int: return _i("EXPERT_STACK_CACHE", 0) +def async_prefetch() -> bool: return _b("ASYNC_PREFETCH", "0") +def prefetch_buffer_experts() -> int: return _i("PREFETCH_BUFFER_EXPERTS", 2048) + + +# ============================ AUTOPIN(专家热度持久化 + 启动预热钉死) ============================ +# 总开关。1(默认)：会话中持续累计每层路由热度并节流落盘({EXPERT_DIR}/pool_usage.json)， +# 引擎启动时按历史热度把每层最热专家预填进常驻池并 pin 住(不参与任何驱逐)，消灭冷启动慢热。 +# 0：整个特性零行为(不计数、不读写盘、不预热)。 +def autopin() -> bool: return _b("AUTOPIN", "1") +# 热度文件路径覆盖：缺省 {EXPERT_DIR}/pool_usage.json，结构 {layer: {expert_id: count}}。 +def autopin_usage_file() -> str: return _s("AUTOPIN_USAGE_FILE", "") +# 落盘节流：约每这么多次前向写一次盘(前向边界按「MoE 层号回绕」近似)；atexit 总会兜底写一次。 +# 写同目录临时文件再 rename 防半写损坏；读不到/损坏静默从零开始。 +def autopin_save_every() -> int: return _i("AUTOPIN_SAVE_EVERY", 200) +# 预热 pin 预算占比：每层 pin 数量 = floor(该层池容量 cap_for × FRAC)，取值 0-1，默认 0.5。 +# pinned 永占真实区槽：FRAC 过大会挤掉 demand 自由槽(驱逐候选变少)，不建议 >0.5。 +def autopin_budget_frac() -> float: return min(1.0, max(0.0, _f("AUTOPIN_BUDGET_FRAC", 0.5))) + + +# ============================ 自定义 Metal 算子 ============================ +def custom_qproj() -> bool: return _b("CUSTOM_QPROJ", "0") +def custom_qproj_bits() -> int: return _i("CUSTOM_QPROJ_BITS", 6) +def custom_qproj_targets() -> str: return _s("CUSTOM_QPROJ_TARGETS", "gate,up") +def custom_qproj_max_seq() -> int: return _i("CUSTOM_QPROJ_MAX_SEQ", 4) +def custom_qproj_tile() -> int: return _i("CUSTOM_QPROJ_TILE", 4) +def custom_fused_moe() -> bool: return _b("CUSTOM_FUSED_MOE", "0") +def custom_fused_moe_bits() -> int: return _i("CUSTOM_FUSED_MOE_BITS", 6) +def custom_fused_moe_lanes() -> int: return _i("CUSTOM_FUSED_MOE_LANES", 8) +def custom_fused_moe_block() -> int: return _i("CUSTOM_FUSED_MOE_BLOCK", 256) +def custom_fused_moe_max_seq() -> int: return _i("CUSTOM_FUSED_MOE_MAX_SEQ", 4) + + +# ============================ native MoE 后端 ============================ +def native_moe() -> bool: return _b("NATIVE_MOE", "0") +def native_moe_mlx_op() -> bool: return _b("NATIVE_MOE_MLX_OP", "0") +def native_moe_synthetic() -> bool: return _b("NATIVE_MOE_SYNTHETIC", "0") +def native_moe_slot_pool() -> bool: return _b("NATIVE_MOE_SLOT_POOL", "0") +def native_moe_raise() -> bool: return _b("NATIVE_MOE_RAISE", "0") +def native_moe_slot_cap() -> int: return _i("NATIVE_MOE_SLOT_CAP", 96) +def native_moe_stage_cache() -> bool: return _b("NATIVE_MOE_STAGE_CACHE", "1") +def native_moe_stage_cache_experts() -> int: return _i("NATIVE_MOE_STAGE_CACHE_EXPERTS", 96) +def native_moe_stage_bundle_cache() -> bool: return _b("NATIVE_MOE_STAGE_BUNDLE_CACHE", "1") +def native_moe_stage_cache_bundles() -> int: return _i("NATIVE_MOE_STAGE_CACHE_BUNDLES", 16) +def native_moe_stage_prefetch() -> bool: return _b("NATIVE_MOE_STAGE_PREFETCH", "0") + + +# ============================ 流式 blob ============================ +def stream_blob() -> bool: return _b("STREAM_BLOB", "0") +def stream_blob_loader() -> bool: return _b("STREAM_BLOB_LOADER", "0") +def stream_blob_bg() -> bool: return _b("STREAM_BLOB_BG", "0") +def stream_blob_workers() -> int: return _i("STREAM_BLOB_WORKERS", 8) +def stream_blob_window() -> int: return _i("STREAM_BLOB_WINDOW", 3) +def stream_blob_nocache(default: str = "1") -> bool: return _b("STREAM_BLOB_NOCACHE", default) +# 默认 12:双源侧区(cap32+侧区32)K=3 MTP 甜点扫描结果(REPEAT=2)。瓶颈是 SSD 读盘时序而非容量, +# 加大 budget 只会灌满共享读队列→预取到得更晚(timing miss↑)。W=24/B=12 在 hit(0.901)/disk(3874)/ +# tok·s(11.61)/确定性(nmis=2)全面优于旧 W=32/B=16(10.0 tok·s、nmis=43),tok·s +16%。 +def stream_blob_bg_budget(default: int = 12) -> int: return _i("STREAM_BLOB_BG_BUDGET", default) +# 按需 miss 批量并行读：把本层所有 miss 专家收成一批，一次 load_experts(8-worker 并行 pread)， +# 取代逐专家串行 pread（实测串行 6.4GB/s vs 并行 8 路 22GB/s）。仅影响读取调度，结果等价。 +# 默认开：剖析实测 decode 时间 ~92% 花在每层 host 回退路径，其中字节物化是大头；批量堆叠 +# (每段一次 mx.array 取代逐专家 6N 次构造)在 cap=12 把 decode 从 ~466ms 降到 ~408ms(~12%)。 +def batch_miss_read() -> bool: return _b("BATCH_MISS_READ", "1") +# demand miss 走 C++ 原生物化(load_experts_native)：blob_load 把字节 pread 直进 MLX 数组、 +# eval 时在 C++ 跑、绕 GIL、惰性。比 stacked 再快 ~6%(把物化从 acquire 挪到 eval)。 +# 默认关(opt-in)：实测它偶发把末尾 token 推进双稳态慢挡(2-3s/token 的"悬崖")， +# 其 lazy 物化推高内存压力,稳定性不如 stacked(纯 numpy 同步、4 次跑 0 慢挡)。仅在确认 +# 环境内存充裕、需要那 6% 时手动开。不消费 async prefetch buffer，故仅在 not async_prefetch 接入。 +def native_demand_loader() -> bool: return _b("NATIVE_DEMAND_LOADER", "0") +def staging_ring() -> int: return _i("STAGING_RING", 2) # 安全下限=2：MTP 每步 verify+replay 各对同层 submit 一次 +def staging_pread_parallel() -> bool: return _b("STAGING_PREAD_PARALLEL", "0") # staging fill 派后台池并行;默认关:实测不降 timing miss(IO 受限)且弱化 buffer 新鲜度不变量,详见 benchmarks/reports/staging-pread-parallel-2026-06-25.md +def stg_verify() -> bool: return _b("STG_VERIFY", "0") # 诊断:acquire_gpu 命中后池槽字节真值校验,默认关、对主路径零影响 +def stream_blob_prefetch_budget(default: int) -> int: return _i("STREAM_BLOB_PREFETCH_BUDGET", default) + + +# ============================ 跨层 / 预取 ============================ +# 注意 AHEAD 默认按调用方传：跨层 hook=0、native 预取=1。 +def cross_layer_prefetch() -> bool: return _b("CROSS_LAYER_PREFETCH", "0") +def cross_layer_ahead(default: int = 0) -> int: return _i("CROSS_LAYER_PREFETCH_AHEAD", default) +def cross_layer_mult() -> int: return max(1, _i("CROSS_LAYER_PREFETCH_MULT", 2)) +# 预测宽度（方案B）：预测 top-N 候选用于"减常驻"，N 大→recall 高且不占内存（只是 gate argpartition）。 +# 真正占 staging 的是过滤常驻后、按分截断到 staging budget 的缺口子集。 +# 默认 24:双源侧区 K=3 甜点扫描(REPEAT=2)。realized hit 在 W≥24 就到顶(~0.90),再加宽只抬"名义 +# recall"——多出的候选在 SSD 时序上落不进侧区且抢带宽,W=48→hit 0.85、W=64→0.75、tok·s 一路下滑。 +# 收窄到刚好覆盖可及工作集(W=24)最快最省。详见 benchmarks/reports/prefetch-width-budget-sweep-2026-07-01.md。 +def cross_layer_predict_width() -> int: return _i("CROSS_LAYER_PREDICT_WIDTH", 24) +# per-layer 自适应 ahead：早层用小 ahead 保召回、晚层用大 ahead 保时序（默认来自实测）。 +def cross_layer_cutoff() -> int: return _i("CROSS_LAYER_CUTOFF", 6) # 切点：层号 int: return _i("CROSS_LAYER_AHEAD_LO", 1) # 早层 ahead（保召回） +def cross_layer_ahead_hi() -> int: return _i("CROSS_LAYER_AHEAD_HI", 3) # 晚层 ahead（保时序） +def predict_use_x() -> bool: return _b("PREDICT_USE_X", "1") # 默认用本层 MoE 输入 x（更新鲜，+3.6pp recall）；=0 回退旧 norm 路径 +def predict_agg() -> str: return _s("PREDICT_AGG", "max") # K+1 token 聚合：max|mean|union +def predict_union_k() -> int: return _i("PREDICT_UNION_K", 8) # union 时每 token 取的 top-k（控候选数） +def native_fused_prefetch() -> bool: return _b("NATIVE_FUSED_PREFETCH", "0") +# 池侧区零拷贝双源(单/双缓冲)：opt-in、默认 off。VirtualPool 收口，消掉 promote 拷贝。 +# 侧区有两种淘汰策略(SIDEREGION_LFU 门控): +# - 旧"∉P 全清"(默认):侧区=一次性预取批,不积累→hit 仅 0.709(反低于基线 0.763)。 +# - 新"持久 LFU"(SIDEREGION_LFU=1,单代 spec_gens=1):跨步累积热专家,只读新增。 +# 实测(80B,cap=32,warmup64,见 report sideregion-lfu-2026-07-01): +# LFU spec=8 → hit 0.73 / active 4.76GB(省内存)；LFU spec=32 → hit 0.81 / 6.9GB(提命中,+8% tok/s)。 +# 命中在 ~0.81 饱和(加 warmup 无效),0.85+ 仍需真加常驻槽(cap=64→0.869)。 +# 注:dual on 各 spec 均有 run-to-run token 漂移(良性时序噪声,字节校验 0 BAD),故默认 off。 +def zerocopy_dual_source() -> bool: return _b("ZEROCOPY_DUAL_SOURCE") +def pool_spec_slots() -> int: return _i("POOL_SPEC_SLOTS", 3) # 每层侧区投机槽数(LFU 推荐 8 省内存 / 32 提命中) +def sideregion_lfu() -> bool: return _b("SIDEREGION_LFU", "1") # 侧区持久 LFU 单缓冲二级缓存(默认 on=生产路径);SIDEREGION_LFU=0 回退 legacy 双缓冲 +def native_no_submit() -> bool: return _b("NATIVE_NO_SUBMIT", "0") +def native_no_promote() -> bool: return _b("NATIVE_NO_PROMOTE", "0") +def native_materialize() -> bool: return _b("NATIVE_MATERIALIZE", "0") +def file_prefetch_global_budget() -> int: return _i("FILE_PREFETCH_GLOBAL_BUDGET", 64) +def stage_prefetch_global_budget() -> int: return _i("STAGE_PREFETCH_GLOBAL_BUDGET", 64) +def stage_prefetch_min_score() -> float: return _f("STAGE_PREFETCH_MIN_SCORE", 0) +def stage_prefetch_per_layer_budget(default: int = 4) -> int: return _i("STAGE_PREFETCH_PER_LAYER_BUDGET", default) + + +# ============================ 前缀快照(跨会话 system+tools 复用) ============================ +# 每个会话开头都是同一段长 system+tools 提示词(agentic 场景 ~2600 tok),新会话首轮 +# 要全量 prefill(实测 158s)。快照前 N 个 token 的 KV/递归态,后续同前缀会话直接跳过。 +# 0=关闭。注意:仅当 prompt 明显长于该值才启用(短提示词无收益)。 +def prefix_snapshot_head() -> int: return _i("PREFIX_SNAPSHOT_HEAD", 2048) +# 快照条目上限(LRU),每条约占几十 MB(线性层递归态为主) +def prefix_snapshot_entries() -> int: return _i("PREFIX_SNAPSHOT_ENTRIES", 4) + + +# ============================ 工具白名单 ============================ +# 渲染进 chat template 的工具白名单(按工具名),用于把 prompt 头部钉成固定的一段。 +# +# 快照命中的前提是头部**字节级**相同(key = ids[:head] 的哈希)。而调用方的工具发现是动态的: +# 同一段 system 提示词会配不同的工具集发来,头部随之变化,快照必然 miss —— 每个新会话都要 +# 全量 prefill 4000 token 上下。实测本部署已因此存下两份不同头部的快照(3972 / 4261 token), +# 互不命中。白名单把工具集收敛成固定集合后,头部才稳定、快照才能真正消灭"首轮等待"。 +# 省 token 只是附带好处(工具定义约占头部三分之一)。 +# +# 默认只放画像(profile)子系统实际在用的 4 个工具;其余(如排班 metro_*)本部署不处理。 +# 设为空(TOOLS_ALLOW=)即关闭过滤、原样透传全部工具。 +_TOOLS_ALLOW_DEFAULT = ("profile_list_employees,profile_get_portrait_overview," + "profile_get_person_report,profile_get_team_report") + + +def tools_allow() -> "list[str]": + spec = _s("TOOLS_ALLOW", _TOOLS_ALLOW_DEFAULT) + return [s for s in (part.strip() for part in spec.split(",")) if s] + + +# ============================ MTP ============================ +def mtp_verify_mode() -> str: return _s("MTP_VERIFY_MODE", "batch") +# MTP drafter 量化位宽(load_mtp bits 参数):4(默认,省内存)| 8(草稿更准,接受率更高, +# 实测杠杆:accept_len 提升即等比例提速;代价 ~0.7GB 内存) +def mtp_bits() -> int: return _i("MTP_BITS", 4) +# 喂给 MTP drafter 的主模型 hidden:pre_norm(默认,与训练/验证一致,接受率高)| post_norm(旧行为) +def mtp_hidden() -> str: return _s("MTP_HIDDEN", "pre_norm") +# 分块 prefill 的每块 token 数:整段 prefill 一次前向的激活峰值 ∝ prompt 长度(每 MoE 层瞬时 +# 物化大量唯一专家 + 长序列激活),把 prompt 切成小块逐块喂入(KV/SSM 按 offset 因果累积、末块 +# 末位 logits 与整段等价),峰值压回 ∝chunk,使 prefill 与 decode 同稳态。默认 2(沿用 DeepSeek +# 实证:每块唯一专家少、稳走 resident 池便宜路径);PREFILL_CHUNK=0 关闭分块(整段 prefill)。 +def prefill_chunk() -> int: return _i("PREFILL_CHUNK", 2) + + +# ============================ KV 量化(IsoQuant K4/V3)============================ +# 只作用于 12 个全注意力层。SO(4) 块旋转去相关 + 非对称 K4/V3 仿射量化:128k KV 3.0→~0.68 GiB。 +# 线性层(Gated DeltaNet 递归态)不动。质量验收:token 一致率≥95% + logits cosine≥0.99。 +def kv_quant() -> bool: return _b("KV_QUANT", "0") +def kv_k_bits() -> int: return _i("KV_K_BITS", 4) # Key 位宽(默认 4) +def kv_v_bits() -> int: return _i("KV_V_BITS", 3) # Value 位宽(默认 3,非对称) +def kv_group_size() -> int: return _i("KV_GROUP_SIZE", 64) # 仿射量化分组(整除 head_dim=256) +def kv_rotate() -> bool: return _b("KV_ROTATE", "1") # SO(4) 块旋转去相关(默认开) +def kv_rot_seed() -> int: return _i("KV_ROT_SEED", 0) # 旋转随机种子(data-oblivious、固定) + + +# ============================ profiling / 诊断 ============================ +def window_prof() -> bool: return _b("WINDOW_PROF", "0") +def predict_recall_prof() -> bool: return _b("PREDICT_RECALL_PROF", "0") +def miss_attrib() -> bool: return _b("MISS_ATTRIB", "0") # miss 归因：A(预测到没到位)/B(没预测到) +def route_trace_enabled() -> bool: return _b("ROUTE_TRACE", "0") +def stream_prof() -> bool: return _b("STREAM_PROF", "0") +def probe_predict_only() -> bool: return _b("PROBE_PREDICT_ONLY", "0") +def probe_perlayer_sync() -> bool: return _b("PROBE_PERLAYER_SYNC", "0") +# 预取 host 墙钟探针：量 predict/submit/promote 各段主线程不可重叠的 CPU 时间。默认关、零开销。 +def prefetch_tprof() -> bool: return _b("PREFETCH_TPROF", "0") +# 并集专家数探针:按前向 seq 分桶记每层路由专家并集大小(seq=K 即 MTP verify 的专家并集)。默认关。 +def union_prof() -> bool: return _b("UNION_PROF", "0") +# 接受率 top-k 覆盖探针(>0 且 profile 时):记每个草稿位置 MTP 的 top-k 候选,量模型真实 token +# 是否落在 top-2/top-3 里 = 树形展开的救回上界。默认 0=关。 +def accept_topk() -> int: return _i("ACCEPT_TOPK", 0) +# 最小树:位置1 展开 top-2。仅当 A 链首草稿被拒且 B 候选=真实 token 时,额外跑一次 B 链前向救回。 +# 两分支各为独立 batch=1 seq=K 前向(线性层不能批处理树,故不拍平),per-forward union 不变。默认关。 +def tree_top2() -> bool: return _b("TREE_TOP2", "0") +# 第2 草稿位置(pos1)top-2 救回:在 tree_top2 基础上,额外抽 chainC(第2 位次选分支),当第1 位 +# 命中但第2 位被拒、且 chainC 第2 token=模型真值时改验 chainC。探针实测 pos1「首选错次选对」比例 +# (~11%)高于 pos0(~7%),消融证实接受长度确实 +2.74%(bit-lossless);但每次救回要多一次主模型 +# 前向,该成本在当前硬件/批量下恰好抵消 token 收益,净 tok/s 无提升(见 benchmarks/_bench_p1_ablation)。 +# 故默认关,保护已验证的 pos0 纯路径 tok/s 收益;留作前向变廉价/批量变大时收益翻正的储备,一行 env 开。 +def tree_top2_p1() -> bool: return _b("TREE_TOP2_P1", "0") +# 完整树形验证(batch-of-paths):把 P 条候选路径拍到 batch 维,一次 batched 前向并行验证所有路径, +# 选接受最长的路径提交(提取赢家 row)。每条路径是普通线性序列,故线性层/全注意力层都走成熟的 +# batch 前向(无需改 kernel);单次前向的 batch=P 计算加宽了预取窗口,同时多路径提升接受长度 +# ——一举拿下 accept_len 与 hit_rate 两只鸟。默认关(与 tree_top2/普通链验证互斥,优先级最高)。 +def tree_verify() -> bool: return _b("TREE_VERIFY", "0") +# 树分支数(候选路径条数 P)。当前 drafter 在位置1 展开 top-P,P=2 即 top-2。 +def tree_branches() -> int: return _i("TREE_BRANCHES", 2) +# 置信度门控动态深度(P-MTP 风格):逐位抽草稿时累计置信度 C_i=p0·…·p_i,C_i≥tau 且未到 depth_max +# 就继续加深,否则本步只 verify 到当前深度。低置信步收缩(省专家加载/cap 压力),高置信步至多到 +# depth_max。探针实测置信度对接受率区分度极强(高低置信接受率差 +42~59pp)。消融结论(见 +# benchmarks/reports,_bench_adaptive):收益全部来自"向下收缩",τ=0.3、depth_max=3 即 +5~6% tok/s +# 且 bit-lossless、零额外显存;向上扩到 K=4 反而更慢(第4 位专家加载成本 > 多接受的 token),且在 +# 生产 EXPERT_SLOTS=32 下 seq=4·top_k=40>cap 会溢出致有损。故 depth_max 默认 3(=基础 K,slots=32 安全), +# 要扩到 4 必须同时把 EXPERT_SLOTS 提到 ≥40。默认关,作可选加速路径(与 tree_top2 互斥)。 +def adaptive_depth() -> bool: return _b("MTP_ADAPTIVE_DEPTH", "0") +def conf_tau() -> float: return _f("MTP_CONF_TAU", 0.3) +def depth_max() -> int: return _i("MTP_DEPTH_MAX", 3) +# 合并路径:在动态深度基础上,对"被保留成深链(n>=2)"的步叠加 pos0 top-2 救回(首 token 被拒且 +# top-2=模型真值时改验 B 链)。两机理正交(动态深度压每步成本、救回抬接受长度),但都作用于低置信步、 +# 方向相反(depth=1 的浅步无位置可救),故叠加非简单相加,须实测。默认关。 +def adaptive_rescue() -> bool: return _b("MTP_ADAPTIVE_RESCUE", "0") + + +def parse_layers_env(name: str) -> "set[int] | None": + """解析 "0-3,5,8" 形式的层集合环境变量；为空返回 None（表示全部层）。""" + spec = os.environ.get(name, "").strip() + if not spec: + return None + out: "set[int]" = set() + for part in spec.split(","): + part = part.strip() + if not part: + continue + if "-" in part: + a, b = part.split("-", 1) + out.update(range(int(a), int(b) + 1)) + else: + out.add(int(part)) + return out diff --git a/mlx_streaming/core/__init__.py b/mlx_streaming/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mlx_streaming/core/cache/__init__.py b/mlx_streaming/core/cache/__init__.py new file mode 100644 index 0000000..e316e52 --- /dev/null +++ b/mlx_streaming/core/cache/__init__.py @@ -0,0 +1 @@ +"""专家缓存模块：常驻池(resident_pool) / LRU+文件后端存储(expert_store) / blob 流式源(blob_loader)。""" diff --git a/mlx_streaming/core/cache/autopin.py b/mlx_streaming/core/cache/autopin.py new file mode 100644 index 0000000..146813f --- /dev/null +++ b/mlx_streaming/core/cache/autopin.py @@ -0,0 +1,174 @@ +"""AUTOPIN：专家热度持久化 + 启动预热钉死。 + +两部分： +- **热度计数与持久化**：搭已有的物化点零同步计频——block host 路径的 flat、 + acquire_gpu 的 LFU piggyback(全命中)与 miss 回退 flat；dual decode/verify 无 Python + 物化点，由 C++ g_real 的 LFU freq 在落盘时按增量差分合并。节流写 + {EXPERT_DIR}/pool_usage.json(同目录临时文件 + rename 防半写；读不到/损坏静默从零)。 +- **启动预热钉死**：build_streaming_model 末尾按历史热度把每层 top-N 热专家 + (N = floor(cap_for × AUTOPIN_BUDGET_FRAC))经 blob 并行读预填进常驻池并 pin 住 + (不参与任何驱逐)，消灭冷启动慢热。 + +覆盖率口径：prefill/host 路径与 GPU remap miss 回退全计；非 dual 的 decode 全命中 +快路径搭 LFU piggyback 计(inds 已物化，零额外同步)；dual decode/verify 靠 C++ freq +兜底(仅 lfu 策略计频)。EVICT_POLICY≠lfu 时两条 GPU 快路径均不计——绝不为计数给 +decode 热路径新增 GPU→host 同步。 +""" +import atexit +import json +import os +import tempfile +from collections import Counter +from typing import Dict, List + +from mlx_streaming import config + +_USAGE_NAME = "pool_usage.json" + + +def usage_path() -> str: + """热度文件路径：AUTOPIN_USAGE_FILE 覆盖，缺省 {EXPERT_DIR}/pool_usage.json。""" + p = config.autopin_usage_file() + return p or os.path.join(config.expert_dir(), _USAGE_NAME) + + +def load_usage(path: str) -> "Dict[int, Counter]": + """读热度文件；不存在/损坏/结构非法一律静默返回空(从零开始)。""" + try: + with open(path) as f: + raw = json.load(f) + return {int(l): Counter({int(e): int(c) for e, c in d.items()}) + for l, d in raw.items()} + except Exception: + return {} + + +def save_usage(path: str, counts: "Dict[int, Counter]") -> None: + """原子写：先写同目录临时文件再 rename，防半写损坏。失败静默(不阻断推理)。""" + try: + d = os.path.dirname(path) or "." + fd, tmp = tempfile.mkstemp(dir=d, prefix=".pool_usage.", suffix=".tmp") + with os.fdopen(fd, "w") as f: + json.dump({str(l): {str(e): c for e, c in cnt.items()} + for l, cnt in counts.items()}, f) + os.replace(tmp, path) + except Exception: + pass + + +class HeatCounter: + """进程内热度累计器：note() 增量计频，节流 + atexit 落盘。 + + 前向边界用「MoE 层号回绕」近似(与 VirtualPool.begin_forward 同判据)：MoE 块按层号 + 递增被调，出现 layer <= 上次即新前向。路径在创建时快照，避免进程退出时 env 已变 + 导致 atexit 写错位置。 + """ + + def __init__(self, path: str): + self.path = path + self.counts: "Dict[int, Counter]" = {} + self._last_layer = -1 + self._forwards_since_save = 0 + # C++ 累计频次的已合并基线(real_freq_dump 是进程累计值，按差分合并防重复计) + self._cpp_baseline: "Dict[int, Dict[int, int]]" = {} + + def note(self, layer: int, ids: List[int]) -> None: + layer = int(layer) + if layer <= self._last_layer: + self._forwards_since_save += 1 + self._last_layer = layer + c = self.counts.get(layer) + if c is None: + c = self.counts[layer] = Counter() + for e in ids: + c[int(e)] += 1 + if self._forwards_since_save >= config.autopin_save_every(): + self.flush() + + def tick(self) -> None: + """显式前向边界(dual decode 无 Python 计频点，由 begin_forward 驱动周期落盘)。""" + self._forwards_since_save += 1 + if self._forwards_since_save >= config.autopin_save_every(): + self.flush() + + def merge_cpp_freq(self) -> None: + """把 dual decode 的 C++ LFU 累计频次按增量差分合并进来。""" + try: + import mlx_streaming.native_moe_ext as _N + flat = _N.real_freq_dump() + except Exception: + return + cur: "Dict[int, Dict[int, int]]" = {} + for i in range(0, len(flat), 3): + l, e, c = int(flat[i]), int(flat[i + 1]), int(flat[i + 2]) + cur.setdefault(l, {})[e] = c + for l, em in cur.items(): + base = self._cpp_baseline.get(l, {}) + cnt = self.counts.setdefault(l, Counter()) + for e, c in em.items(): + b = base.get(e, 0) + cnt[e] += c - b if c >= b else c # real_reset 后 c None: + self.merge_cpp_freq() + save_usage(self.path, self.counts) + self._forwards_since_save = 0 + + +_counter: "HeatCounter | None" = None + + +def _get_counter() -> HeatCounter: + global _counter + if _counter is None: + _counter = HeatCounter(usage_path()) + atexit.register(_counter.flush) + return _counter + + +def note(layer: int, ids: List[int]) -> None: + """路由热度计数入口(热路径调用)：AUTOPIN=0 立即返回，零行为。""" + if not config.autopin(): + return + _get_counter().note(layer, ids) + + +def tick() -> None: + """前向边界入口(virtual_pool begin_forward 调用)：AUTOPIN=0 立即返回。""" + if not config.autopin(): + return + _get_counter().tick() + + +def counter() -> "HeatCounter | None": + """测试/诊断用：当前计数器(未触发过计数则为 None)。""" + return _counter + + +def warm_start_pin(store) -> "dict": + """启动预热钉死：按 usage 文件每层取 top-N 热专家，经 store.pin_batch 预填进池并 pin 住。 + + N = floor(该层 cap_for × AUTOPIN_BUDGET_FRAC)。pin 加载由 pin_batch 收口：有 blob_loader + 走并行 pread(生产唯一数据源)，无则退回逐专家 per-expert 文件(兼容旧部署)。 + usage 缺失/为空 → 返回零摘要(后台计数仍在累计，供下次启动用)；单层失败不阻断整体。 + """ + summary = {"layers": 0, "pinned": 0, "skipped_layers": 0, "seconds": 0.0} + usage = load_usage(usage_path()) + if not usage: + return summary + import time + t0 = time.perf_counter() + frac = config.autopin_budget_frac() + for layer in sorted(usage): + n = int(store.cap_for(layer) * frac) + if n <= 0: + continue + top = [e for e, _ in usage[layer].most_common(n)] + try: + summary["pinned"] += store.pin_batch(layer, top) + summary["layers"] += 1 + except Exception: + summary["skipped_layers"] += 1 + summary["seconds"] = round(time.perf_counter() - t0, 2) + return summary diff --git a/mlx_streaming/core/cache/blob_loader.py b/mlx_streaming/core/cache/blob_loader.py new file mode 100644 index 0000000..8436c05 --- /dev/null +++ b/mlx_streaming/core/cache/blob_loader.py @@ -0,0 +1,238 @@ +"""全流式 blob 专家源：并行 pread + F_NOCACHE + 即时物化为 MLX 量化数组。 + +字节布局由 prep/blob_layout.py 统一描述（单一真相源），支持两种格式： +- v1 affine（expert_blob_v1）：每 proj 段序 [weight, scales, biases]；weight uint32， + scales/biases 存 uint16 原始位、用时 .view(bfloat16)。 +- v2 mxfp4（expert_blob_v2_mxfp4）：每 proj 段序 [weight, scales]（无 biases）； + weight uint32，scales 为 uint8 原始位、绝不 view(bf16)。 +计算复用 MLX quantized_matmul / gather_qmm（不碰已 NO-GO 的 fused kernel）。 +""" +import os +import threading +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor + +import mlx.core as mx +import numpy as np + +from mlx_streaming import config + +# macOS fcntl：F_NOCACHE 提示内核不要把读过的页留在 page cache（支撑低内存目标）。 +_F_NOCACHE = 48 + + +class BlobExpertSource: + def __init__(self, blob_dir: str, hidden: int, inter: int, group: int, bits: int, + num_experts: int, workers: int = 8, nocache: bool = True, + quant_mode: str = "affine", blob_format: "str | None" = None): + self.dir = blob_dir + self.h, self.i, self.g, self.b, self.ne = hidden, inter, group, bits, num_experts + self.workers = workers + self.nocache = nocache + # blob 格式：mxfp4 默认 v2（无 biases、scales uint8），其余默认 v1 affine。 + self.quant_mode = quant_mode + self.blob_format = blob_format or ( + "expert_blob_v2_mxfp4" if quant_mode == "mxfp4" else "expert_blob_v1") + self._segs, self.stride = self._layout() + self._fds: "dict[int, int]" = {} + # 后台预取：只预读原始字节进 _pf_cache（不建 mx.array，可后台线程跑）； + # 滚动窗口只保留最近 window 层的预取字节，支撑低内存。 + self.window = config.stream_blob_window() + self._pf_pool = ThreadPoolExecutor(max_workers=workers) + self._pf_cache: "OrderedDict[tuple, bytes]" = OrderedDict() + self._pf_futures: "dict[tuple, object]" = {} + self._pf_layers: "list[int]" = [] + self._lock = threading.Lock() + self.preads = 0 + self.prefetch_hits = 0 + + def _layout(self): + # 段表统一由 blob_layout.layout_for 给出（dtype 为字符串 "uint32"/"uint16"/"uint8"）。 + from mlx_streaming.prep.blob_layout import layout_for + return layout_for(self.blob_format, self.h, self.i, self.b, self.g) + + def _fd(self, layer: int) -> int: + fd = self._fds.get(layer) + if fd is None: + fd = os.open(os.path.join(self.dir, f"layer{layer:02d}.blob"), os.O_RDONLY) + if self.nocache: + try: + import fcntl + fcntl.fcntl(fd, _F_NOCACHE, 1) + except (OSError, ValueError): + pass + self._fds[layer] = fd + return fd + + def _materialize(self, raw: bytes, view_bf16: bool = True) -> dict: + # view_bf16=False：scales/biases 保留 uint16（.view(bfloat16) 在后台线程会报 + # no Stream(gpu,0)，故后台物化时关掉，由主线程消费时再 view）。 + # 段表 dtype 为字符串：仅 v1 affine 的 uint16 scales/biases 是 bf16 位重解释； + # mxfp4 的 uint8 scales 保持原样，绝不 view(bf16)。 + out = {} + off = 0 + for proj, tensor, dt, shape, nb in self._segs: + np_dt = np.dtype(dt) + view = np.frombuffer(raw, dtype=np_dt, count=nb // np_dt.itemsize, offset=off).reshape(shape) + arr = mx.array(view) + if view_bf16 and dt == "uint16" and tensor in ("scales", "biases"): + arr = arr.view(mx.bfloat16) + out[f"{proj}.{tensor}"] = arr + off += nb + return out + + def _pread(self, layer: int, e: int) -> bytes: + return os.pread(self._fd(layer), self.stride, e * self.stride) + + def prefetch_async(self, layer: int, expert_ids) -> None: + """后台并行预读预测专家的字节（仅 IO，不建 mx.array）。供跨层预测调用。""" + for e in (int(x) for x in expert_ids): + key = (layer, e) + with self._lock: + if key in self._pf_cache or key in self._pf_futures: + continue + self._pf_futures[key] = self._pf_pool.submit(self._pf_job, layer, e) + self._evict_window(layer) + + def _pf_job(self, layer: int, e: int) -> bytes: + raw = self._pread(layer, e) + with self._lock: + self._pf_cache[(layer, e)] = raw + self._pf_futures.pop((layer, e), None) + return raw + + def _evict_window(self, layer: int) -> None: + with self._lock: + if layer not in self._pf_layers: + self._pf_layers.append(layer) + while len(self._pf_layers) > self.window: + old = self._pf_layers.pop(0) + for key in [k for k in self._pf_cache if k[0] == old]: + del self._pf_cache[key] + + def wait_prefetch(self) -> None: + with self._lock: + futs = list(self._pf_futures.values()) + for f in futs: + f.result() + + def read_raw(self, layer: int, expert_ids) -> "list[bytes]": + """取原始字节：优先用预取缓存/在途结果，未命中的并行 pread（仅 IO）。""" + ids = [int(e) for e in expert_ids] + results: "dict[int, bytes]" = {} + misses = [] + for e in ids: + key = (layer, e) + with self._lock: + if key in self._pf_cache: + results[e] = self._pf_cache[key] + self.prefetch_hits += 1 + continue + fut = self._pf_futures.get(key) + if fut is not None: + results[e] = fut.result() + self.prefetch_hits += 1 + else: + misses.append(e) + + def rd(e): + return e, self._pread(layer, e) + + if misses: + self.preads += len(misses) + if self.workers <= 1 or len(misses) <= 1: + for e in misses: + results[e] = self._pread(layer, e) + else: + # 复用持久线程池（之前每次调用新建/销毁 ThreadPoolExecutor → 大量线程创建 + GIL 争用）。 + for e, raw in self._pf_pool.map(rd, misses): + results[e] = raw + return [results[e] for e in ids] + + def load_experts(self, layer: int, expert_ids, view_bf16: bool = True) -> dict: + """返回 {expert_id: {proj.tensor: mx.array}}。并行读字节，lazy 物化 mx.array。 + + view_bf16=False：scales/biases 留 uint16（供后台线程用，避免 .view 报 no Stream）。 + 注：零拷贝(preadv 进 MLX buffer)实测在真实 lazy-eval 流里更慢——为拿可写 buffer + 需 per-expert mx.eval 同步，代价高于省掉的拷贝(见报告)。故保留 lazy frombuffer 路径。 + """ + ids = [int(e) for e in expert_ids] + raws = self.read_raw(layer, ids) + return {e: self._materialize(raw, view_bf16=view_bf16) for e, raw in zip(ids, raws)} + + def load_experts_stacked(self, layer: int, expert_ids, view_bf16: bool = True) -> dict: + """批量物化：每段只建一个 (N,*shape) mx.array（np.stack 后单次构造）， + 取代逐专家 6N 次 frombuffer+mx.array。返回 {proj.tensor: (N,*shape)}，按 ids 顺序。 + + 与 load_experts 同源（同段表、同 view_bf16 语义），只是把"逐专家逐段"折成 + "逐段一次堆叠"——demand 多 miss 消费的碎 mx.array 构造从 6N 降到 6。 + """ + ids = [int(e) for e in expert_ids] + raws = self.read_raw(layer, ids) + out = {} + off = 0 + for proj, tensor, dt, shape, nb in self._segs: + np_dt = np.dtype(dt) + cnt = nb // np_dt.itemsize + views = [np.frombuffer(raw, dtype=np_dt, count=cnt, offset=off).reshape(shape) + for raw in raws] + arr = mx.array(np.stack(views, axis=0)) # (N, *shape) 单次构造 + if view_bf16 and dt == "uint16" and tensor in ("scales", "biases"): + arr = arr.view(mx.bfloat16) + out[f"{proj}.{tensor}"] = arr + off += nb + return out + + def load_experts_native(self, layer: int, expert_ids, view_bf16: bool = False) -> dict: + """native 物化：用 C++ blob_load 把字节 pread 进 MLX 数组（eval 时在 C++ 跑、绕 GIL）， + Python 侧只建惰性切片/view 图（无数据拷贝）。目的是把"拷贝"挪出 GIL，减少后台线程 + 对主线程 per-layer 派发的争用。 + """ + import mlx_streaming.native_moe_ext as _N + ids = [int(e) for e in expert_ids] + path = os.path.join(self.dir, f"layer{layer:02d}.blob") + raw = _N.blob_load(path, mx.array(ids, dtype=mx.uint32), self.stride) # [n, stride] uint8（lazy） + out = {} + for i, e in enumerate(ids): + row = raw[i] + d = {} + off = 0 + for proj, tensor, dt, shape, nb in self._segs: + seg = row[off:off + nb] + # 按段 dtype 重解释：weight→uint32；v1 affine 的 uint16 scales/biases + # 可再 view(bf16)；mxfp4 的 uint8 scales 保持 uint8。 + if dt == "uint32": + arr = seg.view(mx.uint32).reshape(shape) + elif dt == "uint16": + arr = seg.view(mx.uint16).reshape(shape) + if view_bf16: + arr = arr.view(mx.bfloat16) + else: # uint8（mxfp4 scales） + arr = seg.reshape(shape) + d[f"{proj}.{tensor}"] = arr + off += nb + out[e] = d + return out + + def keys(self) -> "list[str]": + return [f"{proj}.{tensor}" for proj, tensor, _, _, _ in self._segs] + + def acquire(self, layer: int, expert_ids): + """对齐 ResidentExpertPool.acquire：返回 (pool_arrays, slots)。 + + pool_arrays: {proj.tensor: (n_uniq,...) 堆叠数组}，喂给 _sub.forward。 + slots: 每个 expert_id 在 unique 列表中的下标（= local，reshape 前）。 + """ + ids = [int(e) for e in expert_ids] + uniq = list(dict.fromkeys(ids)) + experts = self.load_experts(layer, uniq) + pool = {k: mx.stack([experts[e][k] for e in uniq], axis=0) for k in self.keys()} + pos = {e: i for i, e in enumerate(uniq)} + slots = [pos[e] for e in ids] + return pool, slots + + def close(self): + self._pf_pool.shutdown(wait=True) + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() diff --git a/mlx_streaming/core/cache/expert_store.py b/mlx_streaming/core/cache/expert_store.py new file mode 100644 index 0000000..45b0781 --- /dev/null +++ b/mlx_streaming/core/cache/expert_store.py @@ -0,0 +1,604 @@ +"""路线 B 专家后端：每层独立 LRU 缓存 + 按需切片，统计命中率。 + +与早期版本的两处提速改动： +- **每层独立 LRU**：每个 MoE 层各自一个 LRU，容量 = 每层槽数。解码自回归时相邻 + token 在同一层常路由到重叠专家，per-layer 隔离能显著抬高命中率（不再被别层挤掉）。 +- **去掉每次驱逐的 clear_cache 抖动**：驱逐只丢引用，靠 MLX 自带缓冲复用 + 外部 + set_cache_limit 控制上限，避免低命中率下疯狂 malloc/free。 + +stacked[layer] = {"weight": (E,O,I) [, "scales", "biases"]}（惰性，未 eval）。 +fetch(layer, expert_ids) 只取这几个专家、堆叠成 (k, ...) 返回。 +""" +import json +import os +import struct +from collections import OrderedDict, Counter +from typing import Dict, List + +import mlx.core as mx +import numpy as np + +from mlx_streaming import config +from mlx_streaming.core.mem import clear_cache +from mlx_streaming.core.cache.resident_pool import ResidentExpertPool, _POOL_INIT_SLOTS # noqa: F401 + +# macOS fcntl：F_NOCACHE 提示内核不要把读过的页留在 page cache。 +_F_NOCACHE = 48 +# safetensors dtype -> (numpy 读入类型, 需 view 的 mlx 类型)。BF16 numpy 无原生类型, +# 先读 uint16 原始位再 .view(bfloat16)。其余类型按需补。 +_ST_DTYPE = { + "U32": (np.uint32, None), + "U16": (np.uint16, None), + "I32": (np.int32, None), + "F32": (np.float32, None), + "F16": (np.float16, None), + "BF16": (np.uint16, mx.bfloat16), +} + + +def _pread_all(fd: int, size: int) -> bytes: + """把整文件读出(os.pread 可能短读,循环补齐)。""" + chunks = [] + off = 0 + while off < size: + b = os.pread(fd, size - off, off) + if not b: + break + chunks.append(b) + off += len(b) + return b"".join(chunks) + + +def load_safetensors_nocache(path: str) -> "Dict[str, mx.array]": + """用 F_NOCACHE 读 safetensors,返回与 mx.load 等价的 {name: mx.array}。 + + 绕过 OS page cache:每次都真实读 NVMe,基准结果可复现(不被页缓存冷热污染)。 + 与 mx.load 的数值/dtype 完全一致(U32 权重、BF16 scales/biases)。 + """ + fd = os.open(path, os.O_RDONLY) + try: + try: + import fcntl + fcntl.fcntl(fd, _F_NOCACHE, 1) + except (OSError, ValueError, ImportError): + pass + size = os.fstat(fd).st_size + data = _pread_all(fd, size) + finally: + os.close(fd) + n = struct.unpack(" "OrderedDict[int, Dict[str, mx.array]]": + c = self._caches.get(layer) + if c is None: + c = OrderedDict() + self._caches[layer] = c + return c + + def resident_count(self) -> int: + return sum(len(c) for c in self._caches.values()) + + def get_or_load(self, layer: int, e: int, loader) -> Dict[str, mx.array]: + cache = self._layer_cache(layer) + if e in cache: + self.hits += 1 + cache.move_to_end(e) + else: + self.misses += 1 + cache[e] = loader(layer, e) + cache.move_to_end(e) + while len(cache) > self.capacity: + cache.popitem(last=False) # 仅丢引用；不在热路径 clear_cache + if self.clear_on_evict: + clear_cache() + return cache[e] + + def hit_rate(self) -> float: + tot = self.hits + self.misses + return self.hits / tot if tot else 0.0 + + +def _stack_picked(picked: List[Dict[str, mx.array]]) -> Dict[str, mx.array]: + keys = picked[0].keys() + return {k: mx.stack([p[k] for p in picked]) for k in keys} + + +def _stack_cache_size(size: int | None) -> int: + if size is not None: + return max(0, int(size)) + return max(0, config.expert_stack_cache()) + + +def _stack_cache_get(cache, size: int, layer: int, key: tuple[int, ...]): + if size <= 0: + return None + layer_cache = cache.get(layer) + if layer_cache is None or key not in layer_cache: + return None + layer_cache.move_to_end(key) + return layer_cache[key] + + +def _stack_cache_put(cache, size: int, layer: int, key: tuple[int, ...], + value: Dict[str, mx.array]) -> None: + if size <= 0: + return + layer_cache = cache.get(layer) + if layer_cache is None: + layer_cache = OrderedDict() + cache[layer] = layer_cache + layer_cache[key] = value + layer_cache.move_to_end(key) + while len(layer_cache) > size: + layer_cache.popitem(last=False) + + +class LruExpertStore: + """内存堆叠后端：从常驻（惰性）堆叠张量按专家切片。容量为「每层槽数」。""" + + def __init__(self, stacked: Dict[int, Dict[str, mx.array]], capacity: int, + clear_on_evict: bool = False, stack_cache_size: int | None = None): + self._stacked = stacked + self.capacity = capacity + self._lru = _PerLayerLru(capacity, clear_on_evict) + self.stack_cache_size = _stack_cache_size(stack_cache_size) + self._stack_cache: "Dict[int, OrderedDict[tuple[int, ...], Dict[str, mx.array]]]" = {} + @property + def hits(self) -> int: + return self._lru.hits + + @property + def misses(self) -> int: + return self._lru.misses + + def resident_count(self) -> int: + return self._lru.resident_count() + + def cap_for(self, layer: int) -> int: + return self.capacity + + def _load_one(self, layer: int, e: int) -> Dict[str, mx.array]: + src = self._stacked[layer] + out = {} + for k, arr in src.items(): + sub = arr[e] # 切第 e 个专家 + mx.eval(sub) # 显式物化这一个专家 + out[k] = sub + return out + + def fetch(self, layer: int, expert_ids: List[int]) -> Dict[str, mx.array]: + key = tuple(int(e) for e in expert_ids) + cached = _stack_cache_get(self._stack_cache, self.stack_cache_size, layer, key) + if cached is not None: + for e in key: + self._lru.get_or_load(layer, e, self._load_one) + return cached + picked = [self._lru.get_or_load(layer, int(e), self._load_one) for e in expert_ids] + stacked = _stack_picked(picked) + _stack_cache_put(self._stack_cache, self.stack_cache_size, layer, key, stacked) + return stacked + + def hit_rate(self) -> float: + return self._lru.hit_rate() + + +class FileExpertStore: + """文件后端专家缓存：从离线拆分的 per-expert safetensors 按需加载，每层独立 LRU。 + + 文件命名：{root}/layer{layer:02d}_expert{e:03d}.safetensors， + 内容为扁平 dict，键形如 "gate_proj.weight"/"gate_proj.scales"/"up_proj.weight"... + capacity 语义为「每层槽数」（worst-case 常驻 ≈ capacity × MoE 层数）。 + """ + + def __init__(self, root: str, capacity: int, clear_on_evict: bool = False, + record: bool = False, stack_cache_size: int | None = None, + layer_caps: "Dict[int, int] | None" = None): + self.root = root + self.capacity = capacity + self._lru = _PerLayerLru(capacity, clear_on_evict) + self.stack_cache_size = _stack_cache_size(stack_cache_size) + self._stack_cache: "Dict[int, OrderedDict[tuple[int, ...], Dict[str, mx.array]]]" = {} + # ③ 热专家常驻：钉住的专家永不进 LRU、永不驱逐 + self._pinned: "Dict[int, Dict[int, Dict[str, mx.array]]]" = {} + self.pinned_hits = 0 + # 可选 per-layer bundle：把同层所有 per-expert safetensors 合并成一个文件， + # 运行时每层只 mx.load 一次，减少大量小文件 open/header parse 开销。 + self.use_bundles = config.expert_bundle() + self.bundle_dir = config.expert_bundle_dir(os.path.join(root, "layer_bundles")) + self.bundle_cache_size = config.expert_bundle_cache() + self._bundle_cache: "OrderedDict[int, Dict[str, mx.array]]" = OrderedDict() + self._bundle_key_cache: "Dict[int, Dict[int, list[tuple[str, str]]]]" = {} + self.bundle_loads = 0 + # 连续常驻池后端（acquire 路径用），与 LRU stack 路径共享 _load_one + # layer_caps：每层独立池容量(profile 驱动的无损省内存)，缺省层用全局 capacity + # demand miss 批量加载接线(三选一,acquire 优先 stacked > batch > 逐专家): + # - native(默认,且非 async_prefetch):C++ blob_load 直进 MLX 数组、绕 GIL、惰性,返回 + # 逐专家 {e:{k}} → 走 _place_experts。native 不消费 async prefetch buffer,故 gate 在 + # not async_prefetch(async 路径改用 _batch_load_resident,它会先吃 buffer)。 + # - batch_miss_read(默认):numpy 批量堆叠,每段一次 mx.array(取代逐专家 6N 次构造)。 + # - 都关:逐专家 loader。 + if config.native_demand_loader() and not config.async_prefetch(): + _batch_loader = self._batch_load_native + _stacked_loader = None + elif config.batch_miss_read(): + _batch_loader = self._batch_load_resident + _stacked_loader = (self._batch_load_stacked + if not config.async_prefetch() else None) + else: + _batch_loader = None + _stacked_loader = None + self._resident = ResidentExpertPool( + capacity, loader=self._load_one_resident, layer_caps=layer_caps, + batch_loader=_batch_loader, stacked_batch_loader=_stacked_loader) + # 激活频率统计（校准阶段用） + self.record = record + self._counts: "Dict[int, Counter]" = {} + self.async_prefetch = config.async_prefetch() + self.prefetch_buffer_size = config.prefetch_buffer_experts() + self._prefetch_buffer: "OrderedDict[tuple[int, int], Dict[str, mx.array]]" = OrderedDict() + self.prefetch_buffer_hits = 0 + self.prefetch_submitted = 0 + self.prefetch_dropped = 0 + # dual-source：decode 走 staging 双源命中计数（不写池的零拷贝命中） + self.staging_hits = 0 + # 可选 blob miss-loader（STREAM_BLOB_LOADER=1，由 model_builder 注入）： + # miss 时用每专家连续 blob 的并行 pread 取代 per-expert safetensors 的 mx.load， + # 复用常驻池 GPU-remap 快路径（命中零编排），小 EXPERT_SLOTS 即低内存。 + self._blob_loader = None + # 可选后台预取器（STREAM_BLOB_BG=1，由 model_builder 注入）：在独立 stream 上 + # 提前物化预测专家，promote_prefetched 把它们写进常驻池槽（主线程、便宜 scatter）。 + self._bg = None + + def cap_for(self, layer: int) -> int: + """该层常驻池容量(profile 指定则用之，否则全局 capacity)。""" + return self._resident.cap_for(layer) + + def resident_experts(self, layer: int) -> set[int]: + return self._resident.resident_experts(layer) + + def resident_lru_scores(self, layer: int) -> dict[int, float]: + return self._resident.resident_lru_scores(layer) + + @property + def hits(self) -> int: + return self._lru.hits + self.pinned_hits + self._resident.hits + + @property + def misses(self) -> int: + return self._lru.misses + self._resident.misses + + def resident_count(self) -> int: + pinned = sum(len(d) for d in self._pinned.values()) + return self._lru.resident_count() + pinned + + def pinned_count(self) -> int: + return sum(len(d) for d in self._pinned.values()) + + def path(self, layer: int, e: int) -> str: + import os + return os.path.join(self.root, f"layer{layer:02d}_expert{e:03d}.safetensors") + + def bundle_path(self, layer: int) -> str: + return os.path.join(self.bundle_dir, f"layer{layer:02d}.safetensors") + + def _load_layer_bundle(self, layer: int) -> Dict[str, mx.array]: + cached = self._bundle_cache.get(layer) + if cached is not None: + self._bundle_cache.move_to_end(layer) + return cached + bundle = mx.load(self.bundle_path(layer)) + self.bundle_loads += 1 + self._bundle_cache[layer] = bundle + self._bundle_cache.move_to_end(layer) + key_map: "Dict[int, list[tuple[str, str]]]" = {} + for full_key in bundle: + if not full_key.startswith("expert"): + continue + head, rest = full_key.split(".", 1) + e = int(head[len("expert"):]) + key_map.setdefault(e, []).append((rest, full_key)) + self._bundle_key_cache[layer] = key_map + while len(self._bundle_cache) > self.bundle_cache_size: + old_layer, _ = self._bundle_cache.popitem(last=False) + self._bundle_key_cache.pop(old_layer, None) + return bundle + + def _raw_load_one(self, layer: int, e: int) -> Dict[str, mx.array]: + if self._blob_loader is not None: + # blob 单专家：1 次连续 pread + 物化，结构与 mx.load 的 per-expert dict 一致。 + return self._blob_loader.load_experts(int(layer), [int(e)])[int(e)] + if self.use_bundles and os.path.exists(self.bundle_path(layer)): + bundle = self._load_layer_bundle(layer) + keys = self._bundle_key_cache.get(layer, {}).get(int(e), []) + return {short: bundle[full] for short, full in keys} + # 默认走 F_NOCACHE 读取(绕过 OS page cache):基准每次真实读盘、可复现。 + # mx.load 用 mmap 会经页缓存,重复跑被缓存命中污染→速度飘。EXPERT_NOCACHE=0 切回 mmap。 + if config.expert_nocache(): + w = load_safetensors_nocache(self.path(layer, e)) + else: + w = mx.load(self.path(layer, e)) + # 默认惰性:不在此 eval,让读盘并入 MLX 异步图、与计算重叠,token 末统一 eval。 + # 旧版每专家强制 mx.eval(w) 会每 token 触发 ~42 次同步、打碎流水线(实测 ~2.4× 慢)。 + # 数值完全等价(惰性求值仍精确)。EAGER_EXPERT_LOAD=1 可恢复旧的逐专家强制物化。 + if config.eager_expert_load(): + mx.eval(w) + return w + + def _load_one(self, layer: int, e: int) -> Dict[str, mx.array]: + return self._raw_load_one(layer, e) + + def _load_one_resident(self, layer: int, e: int) -> Dict[str, mx.array]: + """resident demand loader：优先消费 async prefetch buffer/inflight。""" + key = (int(layer), int(e)) + if self.async_prefetch: + cached = self._prefetch_buffer.pop(key, None) + if cached is not None: + self.prefetch_buffer_hits += 1 + self._resident.misses = max(0, self._resident.misses - 1) + return cached + return self._raw_load_one(layer, e) + + def _batch_load_resident(self, layer: int, ids: List[int]) -> Dict[int, Dict[str, mx.array]]: + """批量 resident demand loader：先消费 async prefetch buffer，剩余 miss 一次并行读。 + + 关键：blob 路径用一次 load_experts(layer, rest)（8-worker 并行 pread），取代逐专家 + 串行 pread（实测串行 6.4GB/s → 并行 8 路 22GB/s）。结果与逐专家加载等价。 + """ + ids = [int(e) for e in ids] + out: "Dict[int, Dict[str, mx.array]]" = {} + rest: List[int] = [] + for e in ids: + if self.async_prefetch: + cached = self._prefetch_buffer.pop((int(layer), e), None) + if cached is not None: + self.prefetch_buffer_hits += 1 + self._resident.misses = max(0, self._resident.misses - 1) + out[e] = cached + continue + rest.append(e) + if rest: + if self._blob_loader is not None: + out.update(self._blob_loader.load_experts(int(layer), rest)) + else: + for e in rest: + out[e] = self._raw_load_one(layer, e) + return out + + def _batch_load_stacked(self, layer: int, ids: List[int]) -> Dict[str, mx.array]: + """批量+预堆叠 demand loader：blob 路径一次 load_experts_stacked（并行 pread + 每段一次 + 物化），返回 {k:(N,*shape)} 按 ids 顺序。无 blob loader 时退化为逐专家物化后堆叠。 + + 仅在 async_prefetch 关闭时接入（不消费 prefetch buffer），保证不漏 buffer 命中。 + """ + ids = [int(e) for e in ids] + if self._blob_loader is not None: + return self._blob_loader.load_experts_stacked(int(layer), ids) + per = {e: self._raw_load_one(layer, e) for e in ids} + keys = list(per[ids[0]].keys()) + return {k: mx.stack([per[e][k] for e in ids], axis=0) for k in keys} + + def _batch_load_native(self, layer: int, ids: List[int]) -> "Dict[int, Dict[str, mx.array]]": + """批量 native demand loader:blob 路径用 C++ blob_load 把整批字节 pread 直进 MLX 数组 + (eval 时在 C++ 跑、绕 GIL、惰性),返回逐专家 {e:{proj.tensor:arr}}(供 _place_experts)。 + + 相比 numpy 路径(frombuffer+stack+mx.array 同步拷贝),把字节物化从 acquire 挪到 token 末 + eval、并绕开 GIL。无 blob loader 时退回逐专家 raw 物化(行为等价)。view_bf16=True 以 + 兼容 affine 量化(uint16 scales/biases→bf16);mxfp4 的 uint8 scales 忽略此参数。 + """ + ids = [int(e) for e in ids] + if self._blob_loader is not None: + return self._blob_loader.load_experts_native(int(layer), ids, view_bf16=True) + return {e: self._raw_load_one(layer, e) for e in ids} + + def note(self, layer: int, expert_ids: List[int]) -> None: + """记录一次激活（校准阶段调用），用于挑热专家。""" + c = self._counts.get(layer) + if c is None: + c = Counter() + self._counts[layer] = c + for e in expert_ids: + c[int(e)] += 1 + + def recorded_layers(self) -> List[int]: + return list(self._counts.keys()) + + def hot(self, layer: int, h: int) -> List[int]: + """返回该层激活最频繁的 h 个专家 id。""" + c = self._counts.get(layer) + if not c: + return [] + return [e for e, _ in c.most_common(h)] + + def pin(self, layer: int, expert_ids: List[int]) -> None: + """把这些专家加载并钉为常驻（不计入 LRU、不可驱逐）。""" + d = self._pinned.setdefault(layer, {}) + for e in expert_ids: + e = int(e) + if e not in d: + d[e] = self._load_one(layer, e) + # 连续 resident pool 是 decode/MTP 热路径；同步预取进去并标记不可驱逐， + # 否则旧 _pinned 只会影响 fetch/stack 路径，无法降低 acquire_gpu 的 miss。 + self._resident.pin_loaded(layer, {int(e): d[int(e)] for e in expert_ids}) + + def pin_batch(self, layer: int, expert_ids: List[int]) -> int: + """批量 pin（AUTOPIN 预热用）：语义同 pin，但加载收口为一次批量读，返回 pin 数。 + + 加载路径：有 blob_loader（生产唯一专家源）走 load_experts 一次 8-worker 并行 pread， + 取代 pin 的逐专家串行 per-expert safetensors（部署里 per-expert 文件已删，串行路径 + 会直接失败）；无 blob_loader 退回逐专家 _load_one（兼容旧 per-expert 部署）。 + dual(native demand)模式：真实区由 C++ g_real 权威，走 pin_loaded_dual 注册 real_pin + + C++ 直写池行，不留存数组副本（避免 pinned 专家双倍内存），fetch 路径不受影响。 + """ + layer = int(layer) + ids = list(dict.fromkeys(int(e) for e in expert_ids)) + if not ids: + return 0 + if getattr(self._resident, "_native_demand", False): + if self._blob_loader is not None: + loaded = self._blob_loader.load_experts(layer, ids) + else: + loaded = {e: self._load_one(layer, e) for e in ids} + return self._resident.pin_loaded_dual(layer, loaded) + d = self._pinned.setdefault(layer, {}) + todo = [e for e in ids if e not in d] + if todo: + if self._blob_loader is not None: + d.update(self._blob_loader.load_experts(layer, todo)) + else: + for e in todo: + d[e] = self._load_one(layer, e) + self._resident.pin_loaded(layer, {e: d[e] for e in ids}) + return len(ids) + + def prefetch(self, layer: int, expert_ids: List[int]) -> None: + """预取到连续 resident pool，不计入 demand miss/hit。""" + if not self.async_prefetch: + self._resident.prefetch(layer, expert_ids) + return + resident = self._resident.resident_experts(layer) + for e in dict.fromkeys(int(x) for x in expert_ids): + key = (int(layer), int(e)) + if e in resident: + self._resident.prefetch_hits += 1 + continue + if key in self._prefetch_buffer: + self._resident.prefetch_hits += 1 + continue + self._prefetch_buffer[key] = self._raw_load_one(int(layer), int(e)) + self._prefetch_buffer.move_to_end(key) + self.prefetch_submitted += 1 + self._resident.prefetch_loads += 1 + while len(self._prefetch_buffer) > self.prefetch_buffer_size: + self._prefetch_buffer.popitem(last=False) + self.prefetch_dropped += 1 + + def promote_prefetched(self, layer: int) -> int: + """把后台预取器中该层所有「已物化」专家写进常驻池槽（主线程、默认 stream）。 + + 同层(AHEAD=0)预取的就是本层将用到的专家，应真正放进池槽以转 miss 为 hit。 + 互保护：置入时把「本批就绪的全部专家」当作 current，使它们彼此不被驱逐， + 只挤掉池里陈旧的非本批专家；池满且无可驱逐者时跳过剩余，留给 demand 回退。 + 返回写入的专家数。 + """ + if self._bg is None: + return 0 + ready = self._bg.take_ready_layer(int(layer)) + self._bg.note_promote(int(layer), len(ready)) + if not ready: + return 0 + self._resident._ensure_layer(int(layer)) + slot_of = self._resident._slot_of[int(layer)] + protect = {int(e) for e in ready} + placed = 0 + for e, d in ready.items(): + e = int(e) + if e in slot_of: + continue + try: + self._resident._place_expert(int(layer), e, d, current=protect) + except ValueError: + # 池满且无非本批可驱逐 → 剩余专家留给 demand 路径 + break + placed += 1 + return placed + + def wait_prefetch(self) -> None: + """测试/诊断用：等待当前所有 async prefetch 完成。""" + return + + def release_fetch_cache(self) -> int: + """释放 fetch/stack 路径（prefill 用）的常驻缓存，返回释放的专家数。 + + decode 单 token 走 acquire_gpu/_resident 池，永不再碰 _lru/_stack_cache； + prefill 结束后这套缓存（每层 ≤cap 个专家 + 堆叠张量）是纯死重，约占 + cap×MoE层数 个专家的内存（实测 cap=16 时 ~9.2GB）。在 prefill→decode + 边界调用一次即可把这块预算还给系统/池。MTP 多 token 会按需重新填充。 + """ + freed = self._lru.resident_count() + self._lru._caches.clear() + self._stack_cache.clear() + clear_cache() + return freed + + def reset_stats(self) -> None: + self._lru.hits = 0 + self._lru.misses = 0 + self.pinned_hits = 0 + self._resident.hits = 0 + self._resident.misses = 0 + self._resident.gpu_fastpath = 0 + self._resident.gpu_fallback = 0 + self._resident.prefetch_hits = 0 + self._resident.prefetch_loads = 0 + self.prefetch_buffer_hits = 0 + self.prefetch_submitted = 0 + self.prefetch_dropped = 0 + self.staging_hits = 0 + + def acquire(self, layer: int, expert_ids: List[int]): + """连续常驻池取专家：返回 (pool_arrays, slots)。命中零拷贝，miss 单槽写。""" + if self.record: + self.note(layer, expert_ids) + return self._resident.acquire(layer, expert_ids) + + def acquire_gpu(self, layer: int, inds: mx.array, num_experts: int): + """GPU 侧 slot 重映射(decode 热路径):命中零 host 往返。record 模式才回 CPU 记账。""" + if self.record: + self.note(layer, [int(i) for i in inds.reshape(-1).tolist()]) + return self._resident.acquire_gpu(layer, inds, num_experts) + + def fetch(self, layer: int, expert_ids: List[int]) -> Dict[str, mx.array]: + if self.record: + self.note(layer, expert_ids) + key = tuple(int(e) for e in expert_ids) + cached = _stack_cache_get(self._stack_cache, self.stack_cache_size, layer, key) + if cached is not None: + pinned = self._pinned.get(layer) + for e in key: + if pinned is not None and e in pinned: + self.pinned_hits += 1 + else: + self._lru.get_or_load(layer, e, self._load_one) + return cached + pinned = self._pinned.get(layer) + picked = [] + for e in expert_ids: + e = int(e) + if pinned is not None and e in pinned: + self.pinned_hits += 1 + picked.append(pinned[e]) + else: + picked.append(self._lru.get_or_load(layer, e, self._load_one)) + stacked = _stack_picked(picked) + _stack_cache_put(self._stack_cache, self.stack_cache_size, layer, key, stacked) + return stacked + + def hit_rate(self) -> float: + tot = self.hits + self.misses + return self.hits / tot if tot else 0.0 diff --git a/mlx_streaming/core/cache/kv_quant_patch.py b/mlx_streaming/core/cache/kv_quant_patch.py new file mode 100644 index 0000000..e514149 --- /dev/null +++ b/mlx_streaming/core/cache/kv_quant_patch.py @@ -0,0 +1,116 @@ +"""把 Qwen3-Next 的 12 个全注意力层就地切换到 K4/V3 旋转量化 KV。 + +做法(隔离、低侵入): +- 子类 `_RotatedQuantAttn` 仅重写 `__call__`:RoPE 后对 q/k 旋转(R_k)、v 旋转(R_v), + 存非对称量化 cache,走 asym_quantized_sdpa,输出做 V 逆旋转复原。保参数树不变 + (只换 `__class__`、用 object.__setattr__ 挂常量,不进 nn.Module 的参数/子模块树)。 +- 覆盖 `model.make_cache`:全注意力层→AsymmetricQuantizedKVCache,线性层→ArraysCache(size=2)。 +""" +from __future__ import annotations + +import mlx.core as mx +from mlx_lm.models.cache import ArraysCache +from mlx_lm.models.qwen3_next import Qwen3NextAttention + +from mlx_streaming.core.cache.quant_kv import ( + AsymmetricQuantizedKVCache, + asym_quantized_sdpa, + build_block_so4, + rotate_last, +) + + +class _RotatedQuantAttn(Qwen3NextAttention): + """与原 Qwen3NextAttention.__call__ 数值等价(高 bit 时),叠加旋转 + 非对称量化。""" + + def __call__(self, x, mask=None, cache=None): + B, L, D = x.shape + + q_proj_output = self.q_proj(x) + queries, gate = mx.split( + q_proj_output.reshape(B, L, self.num_attention_heads, -1), 2, axis=-1) + gate = gate.reshape(B, L, -1) + + keys, values = self.k_proj(x), self.v_proj(x) + + queries = self.q_norm(queries).transpose(0, 2, 1, 3) + keys = self.k_norm( + keys.reshape(B, L, self.num_key_value_heads, -1)).transpose(0, 2, 1, 3) + values = values.reshape( + B, L, self.num_key_value_heads, -1).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + # 旋转去相关(RoPE 之后):q、k 用 R_k(分数自动抵消),v 用 R_v(输出后逆旋转)。 + Rk = self._kvq_Rk + Rv = self._kvq_Rv + if Rk is not None: + queries = rotate_last(queries, Rk) + keys = rotate_last(keys, Rk) + values = rotate_last(values, Rv) + + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + output = asym_quantized_sdpa( + queries, keys, values, scale=self.scale, mask=mask, + group_size=self._kvq_gs, k_bits=self._kvq_kb, v_bits=self._kvq_vb) + else: + # 无 cache(罕见,如纯前向探针):走稠密 SDPA,V 仍在旋转空间,下方逆旋转复原。 + output = mx.fast.scaled_dot_product_attention( + queries, keys, values, scale=self.scale, mask=mask) + + if Rv is not None: + output = rotate_last(output, Rv.T) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output * mx.sigmoid(gate)) + + +def _inner(model): + """取到含 .layers 的内层模块(mlx_lm Model 包了一层 .model)。""" + return model.model if hasattr(model, "model") else model + + +def patch_kv_quant(model, *, group_size=64, k_bits=4, v_bits=3, rotate=True, seed=0): + """就地把所有全注意力层切到旋转 + 非对称量化 KV,并覆盖 make_cache。返回 model。""" + inner = _inner(model) + layers = inner.layers + + head_dim = None + for l in layers: + if not l.is_linear: + head_dim = l.self_attn.head_dim + break + if head_dim is None: + return model # 没有全注意力层,无需 patch + + Rk = build_block_so4(head_dim, seed=seed) if rotate else None + Rv = build_block_so4(head_dim, seed=seed + 1) if rotate else None + + for l in layers: + if l.is_linear: + continue + attn = l.self_attn + # 常量挂载用 object.__setattr__,避免进 nn.Module 参数/子模块树(否则破坏权重 save/树遍历)。 + object.__setattr__(attn, "_kvq_Rk", Rk) + object.__setattr__(attn, "_kvq_Rv", Rv) + object.__setattr__(attn, "_kvq_gs", group_size) + object.__setattr__(attn, "_kvq_kb", k_bits) + object.__setattr__(attn, "_kvq_vb", v_bits) + attn.__class__ = _RotatedQuantAttn + + def make_cache(): + return [ + AsymmetricQuantizedKVCache(group_size, k_bits, v_bits) + if not l.is_linear else ArraysCache(size=2) + for l in layers + ] + + # 覆盖实例方法(plain function,nn.Module.__setattr__ 不会把函数纳入参数树)。 + model.make_cache = make_cache + return model diff --git a/mlx_streaming/core/cache/quant_kv.py b/mlx_streaming/core/cache/quant_kv.py new file mode 100644 index 0000000..ea7a2e6 --- /dev/null +++ b/mlx_streaming/core/cache/quant_kv.py @@ -0,0 +1,223 @@ +"""IsoQuant 风格 K4/V3 非对称量化 KV cache(仅用于全注意力层)。 + +设计要点(详见 docs/superpowers/specs/2026-06-30-isoquant-kv-quant-design.md): +- SO(4) 块对角旋转去相关:head_dim=256 切成 64 个 4D 块,每块由两个单位四元数构造 + 一个 SO(4) 旋转(left·right_conj sandwich)。旋转正交,在注意力分数里自动抵消 + (q、k 同旋转 → q·kᵀ 不变);对 V 旋转后需在 SDPA 输出上做逆旋转复原。 +- 非对称仿射量化:K 用 k_bits(默认 4)、V 用 v_bits(默认 3),分别用 mx.quantize 打包, + attention 走 mx.quantized_matmul 快路径(每次调用可传不同 bits),无需自写 Metal。 + +打包宽度用 dim*bits//32(对 3-bit 必须如此:mlx_lm 自带 8*4//bits 公式对 3-bit 会算错)。 +""" +from __future__ import annotations + +import mlx.core as mx +import numpy as np +from mlx.utils import tree_map, tree_reduce +# 注意:make_mask 用的是 cache 模块里的 create_attention_mask(签名含 offset/return_array), +# 与 base 模块同名但不同签名的那个不可混用。 +from mlx_lm.models.cache import create_attention_mask + + +# ----------------------------- SO(4) 块旋转 ----------------------------- + +def _quat_left(q: np.ndarray) -> np.ndarray: + """单位四元数 q 的左乘矩阵 L(q):把 v 映射到 q*v(Hamilton 积)。""" + w, x, y, z = q + return np.array([ + [w, -x, -y, -z], + [x, w, -z, y], + [y, z, w, -x], + [z, -y, x, w], + ], dtype=np.float32) + + +def _quat_right_conj(q: np.ndarray) -> np.ndarray: + """单位四元数 q 的"右乘其共轭"矩阵 R(conj q):把 v 映射到 v*conj(q)。 + + L(qL)·R(conj qR) 即 SO(4) 的双边旋转 v -> qL * v * conj(qR),覆盖全部 4D 旋转。 + """ + w, x, y, z = q + return np.array([ + [ w, x, y, z], + [-x, w, z, -y], + [-y, -z, w, x], + [-z, y, -x, w], + ], dtype=np.float32) + + +def build_block_so4(head_dim: int, seed: int = 0, blocks_of: int = 4) -> mx.array: + """构造 head_dim×head_dim 的块对角正交矩阵,每 blocks_of(=4)维一个 SO(4) 旋转。 + + data-oblivious:仅由 seed 决定,固定可复现,不依赖数据分布。 + """ + assert head_dim % blocks_of == 0, "head_dim 必须被 blocks_of 整除" + rng = np.random.default_rng(seed) + M = np.zeros((head_dim, head_dim), dtype=np.float32) + for b in range(head_dim // blocks_of): + qL = rng.standard_normal(4).astype(np.float32) + qL /= np.linalg.norm(qL) + qR = rng.standard_normal(4).astype(np.float32) + qR /= np.linalg.norm(qR) + blk = _quat_left(qL) @ _quat_right_conj(qR) # 4×4 ∈ SO(4) + i = b * blocks_of + M[i:i + blocks_of, i:i + blocks_of] = blk + return mx.array(M) + + +def rotate_last(x: mx.array, R: mx.array) -> mx.array: + """对末维做正交旋转:x[..., D] @ R[D, D](按 x 的 dtype 计算)。""" + return x @ R.astype(x.dtype) + + +# --------------------- 非对称量化 KV cache(K4/V3)--------------------- + +def _packed_words(dim: int, bits: int) -> int: + """量化后每行占的 uint32 个数:dim*bits//32(对 2/3/4/5/6/8 bit 均正确)。""" + return dim * bits // 32 + + +class AsymmetricQuantizedKVCache: + """K、V 各用独立位宽的量化 KV cache(仿 mlx_lm.QuantizedKVCache,但 K/V 不同 bits)。 + + keys/values 各存为 (wq, scales, biases) 三元组,与 mx.quantized_matmul 接口对齐。 + """ + + step = 256 + + def __init__(self, group_size: int = 64, k_bits: int = 4, v_bits: int = 3): + self.keys = None + self.values = None + self.offset = 0 + self.group_size = group_size + self.k_bits = k_bits + self.v_bits = v_bits + + def _init_quant(self, dim, bits, B, H, steps, dt): + packed = _packed_words(dim, bits) + ngroups = dim // self.group_size + return ( + mx.zeros((B, H, steps, packed), dtype=mx.uint32), + mx.zeros((B, H, steps, ngroups), dtype=dt), + mx.zeros((B, H, steps, ngroups), dtype=dt), + ) + + def update_and_fetch(self, keys, values): + B, H, n, kD = keys.shape + vD = values.shape[-1] + prev = self.offset + + if self.keys is None or (prev + n) > self.keys[0].shape[-2]: + new_steps = (self.step + n - 1) // self.step * self.step + if self.keys is not None: + if prev % self.step != 0: + self.keys = tree_map(lambda x: x[..., :prev, :], self.keys) + self.values = tree_map(lambda x: x[..., :prev, :], self.values) + + def _expand(x): + z = mx.zeros((B, H, new_steps, x.shape[-1]), dtype=x.dtype) + return mx.concatenate([x, z], axis=-2) + + self.keys = tree_map(_expand, self.keys) + self.values = tree_map(_expand, self.values) + else: + self.keys = self._init_quant(kD, self.k_bits, B, H, new_steps, keys.dtype) + self.values = self._init_quant(vD, self.v_bits, B, H, new_steps, values.dtype) + + self.offset += n + kq = mx.quantize(keys, group_size=self.group_size, bits=self.k_bits) + vq = mx.quantize(values, group_size=self.group_size, bits=self.v_bits) + for i in range(3): + self.keys[i][..., prev:self.offset, :] = kq[i] + self.values[i][..., prev:self.offset, :] = vq[i] + + return ( + tree_map(lambda x: x[..., :self.offset, :], self.keys), + tree_map(lambda x: x[..., :self.offset, :], self.values), + ) + + def make_mask(self, *args, **kwargs): + return create_attention_mask(*args, offset=self.offset, **kwargs) + + def is_trimmable(self): + return True + + def trim(self, n): + n = min(self.offset, n) + self.offset -= n + return n + + def empty(self): + return self.keys is None + + @property + def nbytes(self): + if self.keys is None: + return 0 + return tree_reduce(lambda a, x: a + x.nbytes, (self.keys, self.values), 0) + + # ---- 快照接口(与 KVCache 同口径,供 mtp.kv_cache._snapshot/_restore 与前缀快照使用)---- + @property + def state(self): + # buffer 未满时切到 offset(MLX 切片产生新数组,快照不会被后续原地写污染) + if self.offset == self.keys[0].shape[-2]: + return self.keys, self.values + return ( + tree_map(lambda x: x[..., : self.offset, :], self.keys), + tree_map(lambda x: x[..., : self.offset, :], self.values), + ) + + @state.setter + def state(self, v): + self.keys, self.values = v + self.offset = self.keys[0].shape[-2] + + @property + def meta_state(self): + return str(self.group_size), str(self.k_bits), str(self.v_bits) + + @meta_state.setter + def meta_state(self, v): + self.group_size, self.k_bits, self.v_bits = map(int, v) + + +# ----------------- 非对称量化 SDPA(K 用 k_bits、V 用 v_bits)----------------- + +def asym_quantized_sdpa(queries, q_keys, q_values, scale, mask, + group_size, k_bits, v_bits): + """基于 mlx_lm.quantized_scaled_dot_product_attention,K/V 各自传 bits。 + + queries: [B, n_q_heads, L, D];q_keys/q_values: 量化三元组(n_kv_heads)。 + 返回 [B, n_q_heads, L, D](仍在旋转后的 V 空间,调用方负责逆旋转)。 + """ + B, n_q_heads, L, D = queries.shape + n_kv_heads = q_keys[0].shape[-3] + n_repeats = n_q_heads // n_kv_heads + + queries = queries * scale + + if n_repeats > 1: + queries = mx.reshape(queries, (B, n_kv_heads, n_repeats, L, D)) + q_keys = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_keys) + q_values = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_values) + + scores = mx.quantized_matmul( + queries, *q_keys, transpose=True, group_size=group_size, bits=k_bits) + if mask is not None: + if isinstance(mask, str): + qL, kL = scores.shape[-2:] + q_indices = mx.arange(kL - qL, kL) + k_indices = mx.arange(kL) + mask = q_indices[:, None] >= k_indices[None] + if mask.dtype == mx.bool_: + scores = mx.where(mask, scores, mx.finfo(scores.dtype).min) + else: + scores = scores + mask + scores = mx.softmax(scores, axis=-1, precise=True) + out = mx.quantized_matmul( + scores, *q_values, transpose=False, group_size=group_size, bits=v_bits) + + if n_repeats > 1: + out = mx.reshape(out, (B, n_q_heads, L, D)) + + return out diff --git a/mlx_streaming/core/cache/resident_pool.py b/mlx_streaming/core/cache/resident_pool.py new file mode 100644 index 0000000..75b25a2 --- /dev/null +++ b/mlx_streaming/core/cache/resident_pool.py @@ -0,0 +1,635 @@ +"""常驻专家池：每层一块连续 GPU 张量 + slot LRU，支持按需增长、pin、GPU 侧重映射。 + +设计要点： +- 命中只返回槽位、不写池；miss 只把单个专家原地写进它的槽位（`_write_slot`），避免拷贝整池。 +- 池按需增长（grow-on-demand）：起步 `_POOL_INIT_SLOTS` 行，工作集扩大时 ~1.5× 增至天花板 + `cap_for(layer)` 才开始 LRU 淘汰。默认无 profile 即自动右尺寸，容量内永不超预算。 +- `acquire_gpu`：decode 热路径用 GPU 查找表做纯 GPU slot 重映射，命中层零 host 往返， + 仅一次 miss 标志同步；真 miss 才回退 host 读盘路径。 +""" +import os +from collections import OrderedDict, Counter +from typing import Dict, List + +import mlx.core as mx + +# 诊断门控：staging/acquire_gpu 路径消费侧字节级真值校验（混合 ahead 损坏取证用，默认关）。 +# 在 acquire_gpu 全命中快路径返回前，对本层真实路由命中的每个专家，把池槽字节与磁盘真值 +# 逐 key 比对；不一致即「池槽装错字节」铁证。受 STG_VERIFY=1 控制，对主路径零影响（默认 off）。 +_STG_VERIFY = os.environ.get("STG_VERIFY") == "1" +_stg_verify_state = {"ok": 0, "bad": 0, "printed": 0, "calls": 0, "first_bad_call": None} + +from mlx_streaming import config +from mlx_streaming.core.cache import autopin + +# Route 3 Phase 1 底座：spec/dual 模式下池 buffer 改由 C++ 拥有(mx::allocator + no-op deleter)， +# 地址进程内恒定、永不被 MLX donation/迁移，供侧区/demand 后台 pread 安全直写（替代消费侧 MLX scatter）。 +# POOL_OWNED=0 可强制回退 mx.zeros(仅供 A/B 对照)。 +_POOL_OWNED = os.environ.get("POOL_OWNED", "1") == "1" +# mx.Dtype -> C++ pool_owned_zeros 接受的 dtype 名。 +_DTYPE_NAME = { + mx.uint32: "uint32", mx.uint16: "uint16", mx.uint8: "uint8", + mx.int32: "int32", mx.int16: "int16", + mx.bfloat16: "bfloat16", mx.float16: "float16", mx.float32: "float32", +} + + +def _owned_pool(sample: "Dict[str, mx.array]", n: int) -> "Dict[str, mx.array]": + """用 C++-owned buffer 建 (n,*shape) 的 per-key 池数组（地址恒定，供 C++ 直写）。""" + import mlx_streaming.native_moe_ext as _N + out = {} + for k, v in sample.items(): + name = _DTYPE_NAME.get(v.dtype) + if name is None: + raise RuntimeError(f"pool_owned_zeros 不支持 dtype {v.dtype} (key={k})") + out[k] = _N.pool_owned_zeros([int(n)] + [int(d) for d in v.shape], name) + return out + + +# 池按需增长(grow-on-demand)的初始物理行数。 +# 起步小,工作集扩大时按 ~1.5× 增长,封顶 cap_for(默认=全局 capacity)。 +# 好处:无需 profile 也能自动右尺寸(内存≈实际工作集),且对任何 prompt 自适应、容量内永不超预算。 +_POOL_INIT_SLOTS = 16 + + +class ResidentExpertPool: + """每层一个连续常驻池：(capacity,*shape) 张量 + slot LRU。 + + 命中只返回槽位、不写池；miss 只把单个专家原地写进它的槽位（_write_slot）。 + loader(layer, e) -> Dict[str, mx.array]，单个专家的参数（未堆叠）。 + """ + + def __init__(self, capacity: int, loader, layer_caps: "Dict[int, int] | None" = None, + spec_slots: int = 0, batch_loader=None, stacked_batch_loader=None, + spec_gens: int = 1): + self.capacity = capacity + self.loader = loader + # 可选批量加载器 batch_loader(layer, [ids]) -> {e: expert_dict}：acquire 用它把本层所有 + # miss 一次并行读（8-worker pread），取代逐专家串行 loader。None 时退回串行 loader。 + self.batch_loader = batch_loader + # 可选「批量+预堆叠」加载器 stacked_batch_loader(layer, [ids]) -> {k:(N,*shape)}： + # 在 batch_loader 基础上再把 6N 次碎 mx.array 物化折成每段一次，写池时直接整批 scatter， + # 既省 frombuffer/mx.array 构造又省 _write_slots_batch 里的 mx.stack。优先级最高。 + self.stacked_batch_loader = stacked_batch_loader + self.spec_slots = int(spec_slots) # >0：侧区模式（预分配 cap+spec、禁 grow） + self.spec_gens = max(1, int(spec_gens)) # 侧区代数：双缓冲=2，物理行=cap+spec_gens*spec_slots + # cap_for(layer) 是该层物理行数的「天花板」：profile 指定则用之，否则=全局 capacity。 + # 池按需增长(grow-on-demand)：起步小、随工作集扩大增至天花板才开始 LRU 淘汰。 + # 因此默认(无 profile)即自动右尺寸——内存≈实际工作集,对任意 prompt 自适应, + # 且天花板内永不超预算(超预算只是更慢,不崩)。增长偶发(预热期),稳态零拷贝。 + self.layer_caps: "Dict[int, int]" = dict(layer_caps or {}) + self._pools: "Dict[int, Dict[str, mx.array]]" = {} + self._slot_of: "Dict[int, OrderedDict[int, int]]" = {} + self._free: "Dict[int, list]" = {} + # 每层当前物理已分配行数(按需增长,≤ cap_for)。grow-on-demand 的核心状态。 + self._alloc: "Dict[int, int]" = {} + # 每层 GPU 查找表(全局专家 id → slot,-1=不在池),供 acquire_gpu 在命中层 + # 做纯 GPU 重映射(消除每层 .tolist 同步)。与 _slot_of 同步维护(仅在表已建时)。 + self._slot_table: "Dict[int, mx.array]" = {} + # 每层 pinned 专家集合：预取进 resident pool 后不参与 LRU 驱逐。 + # 用于 K>2 / 小槽位时保住热专家工作集，降低真实 miss。 + self._pinned: "Dict[int, set[int]]" = {} + # 可选替换策略：默认 lfu（短窗口频率 + LRU tie-break，实测比 lru 命中更高）；EVICT_POLICY=lru 回退纯 LRU。 + self.eviction_policy = config.evict_policy().lower() + self.lfu_decay_interval = config.lfu_decay_interval() + self._freq: "Dict[int, Counter[int]]" = {} + self._access_count: "Dict[int, int]" = {} + self.hits = 0 + self.misses = 0 + self.prefetch_hits = 0 + self.prefetch_loads = 0 + # GPU remap 路径取证:命中层走纯 GPU 快路径次数 vs 有 miss 回退 host 次数 + self.gpu_fastpath = 0 + self.gpu_fallback = 0 + # 累计「分不到槽被迫落 0 号槽」的专家数(来自 demand_dual)。任何 >0 都意味着有前向 + # 拿错专家权重算过 → 输出不可信,必须靠调小 PREFILL_CHUNK / 调大槽位消除。 + self.unplaced = 0 + # dual 路径真实区槽状态由 C++ demand_dual 唯一权威(无 opt-out)。spec 模式 + native 已编译即启用; + # 非 spec(spec_slots==0)或 native 缺失时保持 Python 权威路径(仅 prefill/host/非双源用)。 + # 启用后 _slot_of/_free/_freq 在 dual 路径不再维护,resident_experts/_count 改查 C++ g_real(预取过滤要用)。 + self._native_demand = False + if int(spec_slots) > 0: + try: + import mlx_streaming.native_moe_ext as _N + self._native_demand = hasattr(_N, "demand_dual") + except Exception: + self._native_demand = False + if self._native_demand: + # 复刻基线 8-worker 并行读：demand miss 的 pread 派给 BgReader 并行执行（高优队列）。 + import mlx_streaming.native_moe_ext as _N + try: + _N.bg_reader_start(int(os.environ.get("DEMAND_WORKERS", "8")), 0) + except Exception: + pass + if self._native_demand and os.environ.get("DEMAND_TIMING") == "1": + import atexit + import mlx_streaming.native_moe_ext as _N + _N.demand_timing_enable(True) + atexit.register(lambda: print( + "[DEMAND_TIMING ms] inds/pool/side_snap/real_lock/core/build =", + [round(x / 1e3, 1) for x in _N.demand_timings()], flush=True)) + + def cap_for(self, layer: int) -> int: + """该层池容量：profile 指定则用之(上限 capacity)，否则用全局 capacity。""" + return min(self.layer_caps.get(layer, self.capacity), self.capacity) + + def _ensure_layer(self, layer: int): + if layer not in self._slot_of: + self._slot_of[layer] = OrderedDict() + # 空闲槽列表懒填:首次 miss 分配池时再灌入,之后按需增长时追加。 + # 注意:始终原地 mutate 这个 list,绝不重新绑定,确保 acquire 里持有的本地引用同步可见。 + self._free[layer] = [] + self._pinned[layer] = set() + self._freq[layer] = Counter() + self._access_count[layer] = 0 + + def _alloc_pool(self, layer: int, sample: Dict[str, mx.array]): + """首次为某层分配池张量,起步 _POOL_INIT_SLOTS 行(不超过该层天花板)。""" + if self.spec_slots > 0: + n = self.cap_for(layer) + self.spec_gens * self.spec_slots # 预分配满（含 spec_gens 代侧区） + # spec/dual 模式：池由 C++ 拥有，侧区异步直写落此 buffer（Route 3 底座）。 + if _POOL_OWNED: + self._pools[layer] = _owned_pool(sample, n) + else: + self._pools[layer] = { + k: mx.zeros((n,) + v.shape, dtype=v.dtype) for k, v in sample.items()} + else: + n = min(_POOL_INIT_SLOTS, self.cap_for(layer)) + self._pools[layer] = { + k: mx.zeros((n,) + v.shape, dtype=v.dtype) + for k, v in sample.items() + } + self._alloc[layer] = n + real = self.cap_for(layer) if self.spec_slots > 0 else n + self._free[layer].extend(range(real)) # 原地 mutate,不重新绑定；侧区行不进 free，永不被 LRU 分配/驱逐 + + def _grow_pool(self, layer: int, new_n: int): + """把某层池物理行数扩到 new_n(封顶 cap_for),拼接保留已驻行的数据与 slot 索引。""" + if self.spec_slots > 0: + return # 侧区模式：预分配满，永不 grow（保 C++ 写指针稳定） + old_n = self._alloc[layer] + new_n = min(new_n, self.cap_for(layer)) + if new_n <= old_n: + return + old = self._pools[layer] + self._pools[layer] = { + k: mx.concatenate( + [v, mx.zeros((new_n - old_n,) + v.shape[1:], dtype=v.dtype)], + axis=0) + for k, v in old.items() + } + self._free[layer].extend(range(old_n, new_n)) # 新行追加为空闲,原地 mutate + self._alloc[layer] = new_n + + def preallocate(self, layer: int, sample: "Dict[str, mx.array]", cap: int): + """满 cap 预分配 typed 池并 eval 固定 data 指针；幂等。 + + 供"C++ 直写槽"机制使用：调用后该层不再 grow/不再走 _place_expert(MLX scatter)， + 槽位由调用方直接 mutate _slot_of + _set_table 管理、字节由 C++ pread 写入。 + """ + self._ensure_layer(layer) + if layer in self._pools: + return + if _POOL_OWNED: + self._pools[layer] = _owned_pool(sample, cap) # C++ 拥有，地址恒定供直写 + else: + self._pools[layer] = { + k: mx.zeros((cap,) + v.shape, dtype=v.dtype) for k, v in sample.items() + } + mx.eval(list(self._pools[layer].values())) # 固定 data 指针 + self._alloc[layer] = cap + # 登记所有物理行为空闲,供共用分配器(_alloc_slot)按 free 优先分配； + # 已被 slot_of 占用的行不在此列(暖池复用时 slot_of 已满 → free 为空 → 走驱逐)。 + occupied = set(self._slot_of[layer].values()) + self._free[layer].extend(r for r in range(cap) if r not in occupied) + + def allocated_slots(self, layer: int) -> int: + """该层池张量当前物理行数(按需增长,≤ cap_for),用于内存核算/测试。""" + return self._alloc.get(layer, 0) + + def _pool_owned(self, layer: int) -> bool: + """该层池是否为 C++-owned buffer（spec/dual 模式 + POOL_OWNED）。 + owned 池禁止任何 MLX scatter（会重分配 buffer、孤立侧区直写），落池一律走 C++ 直写。""" + return _POOL_OWNED and self.spec_slots > 0 + + def _write_slot(self, layer: int, slot: int, expert: Dict[str, mx.array]): + if self._pool_owned(layer): + self._write_slots_batch(layer, [slot], [expert]) # 单槽也走 C++ 直写，避免 scatter + return + pool = self._pools[layer] + for k, v in expert.items(): + pool[k][slot] = v # de-risk 选定：原地写，单槽、不拷贝整池 + + def _write_slots_batch(self, layer: int, slots: List[int], + experts: "List[Dict[str, mx.array]]") -> None: + """把多个专家一次性写进各自槽位。 + + owned 池：C++ memcpy 直写各行（无 MLX scatter，保 buffer 地址恒定 → 侧区直写不被孤立）。 + 非 owned：每个 key 一次 stack + fancy-index scatter（原碎 kernel 最少化路径）。 + """ + if not slots: + return + pool = self._pools[layer] + if self._pool_owned(layer): + import mlx_streaming.native_moe_ext as _N + keys = list(pool.keys()) + pool_list = [pool[k] for k in keys] + srcs_flat = [e[k] for k in keys for e in experts] # key-major：与 C++ pool_write_rows 约定一致 + mx.eval(srcs_flat) # 物化源，固定 data 指针 + _N.pool_write_rows(pool_list, srcs_flat, [int(s) for s in slots]) + return + idx = mx.array(slots, dtype=mx.int32) + for k in pool: + pool[k][idx] = mx.stack([e[k] for e in experts], axis=0) + + def resident_count(self, layer: int) -> int: + if self._native_demand: # 方案B：C++ g_real 为真实区权威 + import mlx_streaming.native_moe_ext as N + return int(N.real_region_count(int(layer))) + return len(self._slot_of.get(layer, ())) + + def resident_experts(self, layer: int) -> set[int]: + """返回该层 resident pool 里当前已有的专家 id 集合（trace/probe + 侧区预取过滤用）。""" + if self._native_demand: # 方案B：查 C++ g_real，保预取过滤一致 + import mlx_streaming.native_moe_ext as N + flat = N.real_region_contents(int(layer)) + return {int(flat[i]) for i in range(0, len(flat), 2)} + return set(self._slot_of.get(layer, {}).keys()) + + def _bootstrap_dual_pool(self, layer: int) -> None: + """方案B：首次为该层预分配满(cap+spec_gens*spec)池、eval 固定 data 指针 + C++ real_init。幂等。 + + 池结构(per-key typed 数组)由 Python 从一个样本专家的 shape 建；字节后续由 C++ demand pread 写入。 + Python 侧 _slot_of/_free 在方案B dual 路径不再使用（真实区状态归 C++ g_real）。 + """ + self._ensure_layer(layer) + if layer not in self._pools: + sample = self.loader(layer, 0) # 仅取 shape/dtype，不写入池 + cap = self.cap_for(layer) + n = cap + self.spec_gens * self.spec_slots + if _POOL_OWNED: + self._pools[layer] = _owned_pool(sample, n) # C++ 拥有，地址恒定供直写 + else: + self._pools[layer] = { + k: mx.zeros((n,) + v.shape, dtype=v.dtype) for k, v in sample.items() + } + mx.eval(list(self._pools[layer].values())) # 固定 data 指针，供 C++ memcpy + self._alloc[layer] = n + import mlx_streaming.native_moe_ext as N + N.real_init(int(layer), int(cap)) + + def resident_lru_scores(self, layer: int) -> dict[int, float]: + """返回 resident 专家的 LRU 新近度分数：0=最久未用，1=最近使用。""" + keys = list(self._slot_of.get(layer, {}).keys()) + if not keys: + return {} + denom = max(1, len(keys) - 1) + return {int(e): i / denom for i, e in enumerate(keys)} + + def _note_access(self, layer: int, expert_ids: List[int]) -> None: + if self.eviction_policy != "lfu": + return + self._ensure_layer(layer) + freq = self._freq[layer] + for e in expert_ids: + freq[int(e)] += 1 + self._access_count[layer] += len(expert_ids) + if self.lfu_decay_interval > 0 and self._access_count[layer] >= self.lfu_decay_interval: + for e in list(freq): + freq[e] //= 2 + if freq[e] <= 0: + del freq[e] + self._access_count[layer] = 0 + + def _choose_victim(self, layer: int, current: set[int]) -> int: + slot_of = self._slot_of[layer] + pinned = self._pinned.get(layer, set()) + # 绝不驱逐当前请求(current)专家：否则 acquire 末尾按 slot_of 取槽会 KeyError。 + # 只在「非 current 且非 pinned」中选受害者；选不出来说明本次唯一专家+不可驱逐者已超容量， + # 由调用方(acquire 的 len(uniq)>cap 守卫)负责拦截，这里直接报清晰错误。 + candidates = [e for e in slot_of if e not in pinned and e not in current] + if not candidates: + raise ValueError( + f"layer {layer} resident pool has no evictable non-current slot " + f"(capacity={self.cap_for(layer)}, pinned={len(pinned)}, current={len(current)})") + if self.eviction_policy != "lfu": + return candidates[0] + freq = self._freq.get(layer, Counter()) + return min(enumerate(candidates), key=lambda x: (freq.get(x[1], 0), x[0]))[1] + + def _alloc_slot(self, layer: int, e: int, current: "set[int]") -> "tuple[int, bool]": + """为 e 分配/复用槽（free 优先，否则驱逐非 pinned/非 current 的 LRU 最旧）， + 更新 slot_of + table，但不写数据。返回 (slot, is_new)。 + + 供 C++ 直写预取(prefetch_cpp)与 demand 写入(_place_expert)**共用**：保证两条路径 + 从同一套 _free/_slot_of/_slot_table 分配，绝不抢占同一物理槽。 + """ + slot_of, free = self._slot_of[layer], self._free[layer] + if e in slot_of: + slot_of.move_to_end(e) + return slot_of[e], False + cap = self.cap_for(layer) + if not free and self._alloc.get(layer, 0) < cap: + cur = self._alloc[layer] + self._grow_pool(layer, cur + max(1, cur // 2)) + if free: + slot = free.pop(0) + else: + evicted_e = self._choose_victim(layer, current) + slot = slot_of.pop(evicted_e) + self._clear_table(layer, evicted_e) + slot_of[e] = slot + slot_of.move_to_end(e) + self._set_table(layer, e, slot) + return slot, True + + def _place_expert(self, layer: int, e: int, expert: Dict[str, mx.array], + current: "set[int] | None" = None) -> int: + """把专家写进 resident pool，必要时增长或驱逐非 pinned LRU，返回 slot。""" + current = current or {e} + if layer not in self._pools: + self._alloc_pool(layer, expert) + slot, _ = self._alloc_slot(layer, e, current) + self._write_slot(layer, slot, expert) + return slot + + def _place_experts(self, layer: int, ids: List[int], + experts: "Dict[int, Dict[str, mx.array]]", + current: "set[int]") -> List[int]: + """批量把一组 miss 专家写进 resident pool：共用槽分配器 + 单次批量 scatter。 + + 与逐个 _place_expert 语义等价（同样的 grow/驱逐/槽分配顺序）， + 只把数据写入合并成 _write_slots_batch，去掉每 token 上百次碎 scatter。 + """ + if not ids: + return [] + if layer not in self._pools: + self._alloc_pool(layer, experts[ids[0]]) + slots = [self._alloc_slot(layer, e, current)[0] for e in ids] + self._write_slots_batch(layer, slots, [experts[e] for e in ids]) + return slots + + def _place_experts_stacked(self, layer: int, ids: List[int], + stacked: "Dict[str, mx.array]", + current: "set[int]") -> List[int]: + """批量写一组 miss：stacked 为按 ids 顺序预堆叠的 {k:(N,*shape)}， + 每 key 直接一次 fancy-index scatter（连 _write_slots_batch 的 mx.stack 都省了）。 + """ + if not ids: + return [] + if layer not in self._pools: + self._alloc_pool(layer, {k: v[0] for k, v in stacked.items()}) + slots = [self._alloc_slot(layer, e, current)[0] for e in ids] + pool = self._pools[layer] + if self._pool_owned(layer): + import mlx_streaming.native_moe_ext as _N + keys = list(pool.keys()) + pool_list = [pool[k] for k in keys] + stacked_list = [stacked[k] for k in keys] + mx.eval(stacked_list) # 物化预堆叠源 + _N.pool_write_stacked(pool_list, stacked_list, [int(s) for s in slots]) + return slots + idx = mx.array(slots, dtype=mx.int32) + for k in pool: + pool[k][idx] = stacked[k] + return slots + + def prefetch_cpp(self, layer: int, expert_ids, submit_fn) -> list: + """用 C++ 直写把预测专家投机预取进池（与 demand fallback 共用槽分配器）。 + + submit_fn(slot, expert)：提交一次异步 C++ 后台读到该物理槽（调用方负责后续 wait）。 + 已驻 → 触摸 LRU 跳过；未驻 → 分配槽并 submit。返回新提交的 expert 列表。 + 投机性质：不计 hit/miss，分得的槽可被后续 LRU 驱逐；无可驱逐槽时跳过该专家 + （交给 demand fallback），绝不抛错打断热路径。 + """ + self._ensure_layer(layer) + if layer not in self._pools: + return [] # 池未建（首 token 预热）→ 全交 fallback + current = {int(e) for e in expert_ids} + slot_of = self._slot_of[layer] + submitted = [] + for e in expert_ids: + e = int(e) + if e in slot_of: + slot_of.move_to_end(e) + continue + try: + slot, _ = self._alloc_slot(layer, e, current) + except ValueError: + continue # 无可驱逐槽 → 跳过，留给 demand fallback + submit_fn(slot, e) + submitted.append(e) + return submitted + + def pin(self, layer: int, expert_ids: List[int]) -> None: + """预取并钉住专家到 resident pool；pinned 专家不被 LRU 驱逐。 + + pin 是显式预取，不计入 hit/miss 统计；后续 acquire/acquire_gpu 命中才计 hit。 + """ + # dual 模式(demand_dual 唯一权威)已弃用 PIN_HOT:真实区归 C++ g_real,Python pinned 不再生效。 + # 明确报错而非静默无效——请设 PIN_HOT=0。 + if self._native_demand and list(expert_ids): + raise RuntimeError( + "PIN_HOT 在 demand_dual(dual 模式)下已弃用:真实区由 C++ g_real 权威、不支持 pin。" + "请设 PIN_HOT=0。") + loaded = {int(e): self.loader(layer, int(e)) for e in expert_ids} + self.pin_loaded(layer, loaded) + + def pin_loaded(self, layer: int, experts: "Dict[int, Dict[str, mx.array]]") -> None: + """把已加载专家钉进 resident pool，避免 FileExpertStore.pin 重复读/持有两份。""" + uniq = list(dict.fromkeys(int(e) for e in experts)) + cap = self.cap_for(layer) + self._ensure_layer(layer) + pinned = self._pinned[layer] + if len(pinned | set(uniq)) > cap: + raise ValueError( + f"pin 后 pinned 数量 {len(pinned | set(uniq))} > 该层池容量 {cap}") + slot_of = self._slot_of[layer] + for e in uniq: + if e in slot_of: + pinned.add(e) + slot_of.move_to_end(e) + continue + self._place_expert(layer, e, experts[e], current={e}) + pinned.add(e) + + def pin_loaded_dual(self, layer: int, experts: "Dict[int, Dict[str, mx.array]]") -> int: + """dual(native demand)模式的 pin：C++ g_real 是真实区唯一权威（Python _slot_of 不生效）， + 必须 real_pin 注册（pinned 免驱逐 + 从 free 头取真实区槽）再用 pool_write_rows 把字节 + C++ 直写池行（owned 池禁止 MLX scatter）。返回成功 pin 的专家数。 + + 仅在 _native_demand 下由 FileExpertStore.pin_batch 调用；experts 为已加载的 + {expert: {key: arr}}（key 集/顺序与池一致——同源自 blob 段表）。 + """ + import mlx_streaming.native_moe_ext as _N + ids = list(dict.fromkeys(int(e) for e in experts)) + if not ids: + return 0 + self._bootstrap_dual_pool(layer) # 建池 + real_init（幂等） + slots = _N.real_pin(int(layer), ids) # 登记 pinned + 分配真实区槽（-1=分不到） + ok = [(e, s) for e, s in zip(ids, slots) if s >= 0] + if not ok: + return 0 + keys = list(self._pools[layer].keys()) + pool_list = [self._pools[layer][k] for k in keys] + srcs_flat = [experts[e][k] for k in keys for e, _ in ok] # key-major，同 _write_slots_batch + mx.eval(srcs_flat) + _N.pool_write_rows(pool_list, srcs_flat, [s for _, s in ok]) + return len(ok) + + def prefetch(self, layer: int, expert_ids: List[int]) -> None: + """把专家预取进 resident pool，但不计入 hit/miss，且可被后续 LRU 驱逐。""" + uniq = list(dict.fromkeys(int(e) for e in expert_ids)) + cap = self.cap_for(layer) + if len(uniq) > cap: + uniq = uniq[:cap] + self._ensure_layer(layer) + slot_of = self._slot_of[layer] + current = set(uniq) + for e in uniq: + if e in slot_of: + slot_of.move_to_end(e) + self.prefetch_hits += 1 + continue + expert = self.loader(layer, e) + self._place_expert(layer, e, expert, current=current) + self.prefetch_loads += 1 + + def acquire(self, layer: int, expert_ids: List[int], protect: "set[int] | None" = None): + """protect：额外的「不可驱逐」专家集（除本批 miss 外）。dual 路径传入本前向全部 inds， + 确保落 miss 腾槽时不驱逐本前向仍需的命中专家（否则其 slot 被复用 → 消费侧 gather 错字节）。 + 仅影响驱逐候选，不改 hit/miss 记账口径。""" + # 唯一专家集合(保序去重)：池只需同时容纳本次请求的唯一专家 + uniq = list(dict.fromkeys(int(e) for e in expert_ids)) + cap = self.cap_for(layer) + if len(uniq) > cap: + raise ValueError( + f"本次请求 {len(uniq)} 个唯一专家 > 该层池容量 {cap}") + self._ensure_layer(layer) + self._note_access(layer, uniq) + uniq_set = set(uniq) + # 驱逐保护集 = 本批 miss ∪ 调用方额外指定(如本前向全部 inds 命中专家)。 + current = uniq_set if protect is None else (uniq_set | protect) + slot_of, free = self._slot_of[layer], self._free[layer] + # 先处理命中（触摸 LRU），再把所有 miss 收成一批 + misses = [] + for e in uniq: + if e in slot_of: + self.hits += 1 + slot_of.move_to_end(e) # 触摸为最近使用，避免本次内被驱逐 + else: + misses.append(e) + if misses: + self.misses += len(misses) + if self.stacked_batch_loader is not None: + # 批量读 + 批量物化：每段一次 mx.array，写池每 key 一次 scatter（碎 kernel 最少） + stacked = self.stacked_batch_loader(layer, misses) + self._place_experts_stacked(layer, misses, stacked, current=current) + else: + if self.batch_loader is not None: + # 一次并行读本层所有 miss（8-worker pread）；保序写池，LRU/驱逐语义不变 + loaded = self.batch_loader(layer, misses) + else: + loaded = {e: self.loader(layer, e) for e in misses} + # 批量写：每 key 仅一次 stacked scatter，取代逐专家逐段碎 scatter + self._place_experts(layer, misses, loaded, current=current) + # slots 与原始 expert_ids 一一对应(含重复)，便于直接 reshape 成 routing 索引 + slots = [slot_of[int(e)] for e in expert_ids] + return self._pools[layer], slots + + def _set_table(self, layer: int, e: int, slot: int): + t = self._slot_table.get(layer) + if t is not None: + t[int(e)] = slot # (num_experts,) 小张量原地改,代价可忽略 + + def _clear_table(self, layer: int, e: int): + t = self._slot_table.get(layer) + if t is not None: + t[int(e)] = -1 + + def _ensure_table(self, layer: int, num_experts: int) -> mx.array: + t = self._slot_table.get(layer) + if t is None or int(t.shape[0]) != num_experts: + self._ensure_layer(layer) + tab = [-1] * num_experts + for e, slot in self._slot_of[layer].items(): + tab[int(e)] = int(slot) + t = mx.array(tab, dtype=mx.int32) + self._slot_table[layer] = t + return t + + def acquire_gpu(self, layer: int, inds: mx.array, num_experts: int): + """GPU 侧 slot 重映射:命中层零 host 往返(仅一次 miss 标志同步)。 + + 返回 (pool_arrays, local)。inds 为路由结果(decode 时形如 (1,1,k))。 + 全命中 → local = 表[inds] 纯 GPU;有 miss → 回退既有 host acquire 读盘并维护表后重算 local。 + """ + table = self._ensure_table(layer, num_experts) + local = mx.take(table, inds) + n_miss = int(mx.sum((local < 0).astype(mx.int32))) # 唯一一次 GPU→CPU 同步 + if n_miss == 0: + self.gpu_fastpath += 1 + self.hits += int(inds.size) # decode top-k 各专家互异 → size 即唯一命中数 + if self.eviction_policy == "lfu": # 全命中层也计频:piggyback 上面 n_miss 的 drain,inds 已物化 + flat = [int(i) for i in inds.reshape(-1).tolist()] + self._note_access(layer, flat) + autopin.note(layer, flat) # AUTOPIN 热度:搭同一物化点,零额外同步(AUTOPIN=0 立即返回) + return self._pools[layer], local + # 有 miss:回退 host 路径(读盘、写槽、维护表),再用更新后的表重算 local + self.gpu_fallback += 1 + flat = [int(i) for i in inds.reshape(-1).tolist()] + autopin.note(layer, flat) # AUTOPIN 热度:miss 回退 flat 已物化,顺带计频 + pool_arrays, _ = self.acquire(layer, flat) + local = mx.take(self._slot_table[layer], inds) + return pool_arrays, local + + def verify_acquire_bytes(self, layer, inds, stg=None): + """诊断(STG_VERIFY)：acquire 后把本层真实路由命中专家的池槽字节与磁盘真值逐 key 比对。 + + 发现不一致即打印 (call, layer, expert, slot, gen, key)：池槽装错字节的铁证， + 与 timing/gen 竞态对应。call 是全局校验调用序号(≈token 前向序，便于和首分歧 token 关联)。 + stg：可选 NativeStagingManager，用于查该专家落池所用 gen。 + """ + st = _stg_verify_state + st["calls"] += 1 + call = st["calls"] + pool = self._pools.get(layer) + slot_of = self._slot_of.get(layer) + if pool is None or slot_of is None: + return + flat = {int(i) for i in inds.reshape(-1).tolist()} + for e in flat: + slot = slot_of.get(e) + if slot is None: + continue + try: + truth = self.loader(layer, e) + except Exception: + continue + bad_key = None + for k in pool: + if k not in truth: + continue + a = pool[k][slot] + b = truth[k] + if a.shape != b.shape or not bool(mx.all(a == b)): + bad_key = k + break + if bad_key is None: + st["ok"] += 1 + else: + st["bad"] += 1 + if st["first_bad_call"] is None: + st["first_bad_call"] = call + gen = None + if stg is not None: + gen = stg.placed_gen.get((layer, e)) + if st["printed"] < 24: + st["printed"] += 1 + print(f"[STG_VERIFY] BAD call={call} layer={layer} expert={e} " + f"slot={slot} gen={gen} key={bad_key} " + f"(ok={st['ok']} bad={st['bad']})", flush=True) + + def hit_rate(self) -> float: + tot = self.hits + self.misses + return self.hits / tot if tot else 0.0 diff --git a/mlx_streaming/core/cache/virtual_pool.py b/mlx_streaming/core/cache/virtual_pool.py new file mode 100644 index 0000000..d1734fa --- /dev/null +++ b/mlx_streaming/core/cache/virtual_pool.py @@ -0,0 +1,237 @@ +"""VirtualPool：预取调度器 + 双源双缓冲协调器（统一收口）。 + +一个对象承担两职，因为 block.py 用同一个 `self._vpool` 属性： + +1. ahead 调度（两种模式都用）：`ahead_for` / `target_for` 决定「第 L 层应预读哪层」—— + 早层用小 ahead 保召回、cutoff 起用大 ahead 抢时序。`_native_fused_prefetch` 靠它选目标层。 + +2. 双源双缓冲协调（仅 ZEROCOPY_DUAL_SOURCE 模式）：`begin_forward` / `read_gen` / + `fill_gen` / `acquire` / `prefetch`。侧区分两「代(gen)」物理行，每前向读上一代、fill 写 + 另一代，根除「本前向消费 gather 读的物理行在本前向 eval 期被 fill 覆盖」的竞态 + （见 spec 2026-06-25-qwen-virtual-pool-double-buffer）。对外只暴露「专家→物理行」单次 + gather 接口，消费者每层零 host 同步。 + +构造两种签名并存（互斥使用）： +- 调度器：VirtualPool(num_layers=.., cutoff=.., ahead_lo=.., ahead_hi=..) +- 协调器：VirtualPool(resident, staging, spec_slots)，dual-source 下再补调度参数即可两职合一。 +""" +import mlx.core as mx + +from mlx_streaming.core.cache import autopin + +# 方案B STG_VERIFY 校验累计态（诊断用，默认路径不触及）。 +_stg_verify_state = {"ok": 0, "bad": 0, "printed": 0, "calls": 0} + + +class VirtualPool: + def __init__(self, resident=None, staging=None, spec_slots=None, *, + num_layers=None, cutoff=None, ahead_lo=None, ahead_hi=None, store=None): + # --- 双源协调（resident/staging 存在时启用）--- + self._rp = resident + self._stg = staging + self._store = store # acquire_host 超容量 fetch 回退用(blob 感知);调度器形态可为 None + self._spec = int(spec_slots) if spec_slots is not None else 0 + self._gen = 0 + self._last_layer = -1 # 前向边界检测：层号回绕(<= 上次) 即新前向 + # --- ahead 调度 --- + self._num_layers = int(num_layers) if num_layers is not None else 0 + self._cutoff = int(cutoff) if cutoff is not None else 0 + self._a_lo = max(1, int(ahead_lo)) if ahead_lo is not None else 1 + self._a_hi = max(1, int(ahead_hi)) if ahead_hi is not None else 1 + + # ---- ahead 调度 ---- + def ahead_for(self, src_layer: int) -> int: + # cutoff：早层用 lo（保召回），cutoff 起用 hi（保时序）。 + return self._a_lo if int(src_layer) < self._cutoff else self._a_hi + + def target_for(self, src_layer: int) -> int: + # 目标层 = src+ahead。末层无可预读 → 返回 0（跳过）。 + # 不再 clamp 到末层：clamp 会让多个源层同时预读同一末层，对该层每前向 submit 次数 > + # staging ring，其环形 buffer 在惰性切片被本前向 eval 消费前就被后续 submit 完成回调覆盖成 + # 别的专家字节 → 池槽装错字节。越界即跳过（交给 demand 读真值），保证每层每前向至多 1 次 submit。 + L = int(src_layer) + if L >= self._num_layers - 1: + return 0 + tgt = L + self.ahead_for(L) + if tgt > self._num_layers - 1: # 越界（本会触发 clamp 堆叠）→ 跳过 + return 0 + return tgt + + # ---- 双源双缓冲协调 ---- + def begin_forward(self, layer_idx: int): + """每个 MoE 块 __call__ 开头调；层号回绕(<= 上次) 判为新前向 → 代 +1。 + 稳健：不依赖首个 MoE 层是 layer 0、也不要求 MoE 层连续。""" + if layer_idx <= self._last_layer: + self._gen += 1 + # AUTOPIN:dual decode 无 Python 计频点,用前向边界驱动周期落盘(AUTOPIN=0 立即返回)。 + autopin.tick() + # 新前向开头排空上一前向提交的侧区 fill：C++-owned 池 buffer 由后台异步直写， + # 若消费前 fill 未写完，GPU gather 会读到半写侧区行（DUAL_VERIFY BAD）。 + # drain 阻塞到在途 fill 全部写完 → 本前向要消费的侧区行字节必已就绪。 + if self._spec > 0: + import mlx_streaming.native_moe_ext as _N + _N.sideregion_drain() + self._last_layer = layer_idx + + def _gens(self) -> int: + # 代数取自常驻池;单代(=1)时读=填=0(持久 LFU 单区),双代(=2)时交替(%2==&1)。 + g = getattr(self._rp, "spec_gens", 2) if self._rp is not None else 2 + return max(1, int(g)) + + def read_gen(self) -> int: + return (self._gen - 1) % self._gens() # 读上一前向填好的代;单代恒 0 + + def fill_gen(self) -> int: + return self._gen % self._gens() # fill 写本代;单代恒 0 + + def acquire(self, layer, inds, num_experts, *, seq_len=None, layer_cap=None): + """统一取用入口（GPU-remap 路径）：对外呈现「所有专家都在」的视角。 + + 返回 (pool_arrays, local, n_experts)，计算侧零分支： + - dual（有侧区 staging 且 spec>0）：C++ demand_dual 唯一权威（真实区 ∪ 侧区读代单次 gather）； + n_experts = layer_cap + spec_gens*spec_slots。 + - 非 dual GPU-remap：acquire_gpu；n_experts = layer_cap。 + （host/fetch 路径见 acquire_host，其输入是 host 侧已 .tolist 的 flat，语义不同故分开。） + """ + cap = int(layer_cap) if layer_cap is not None else self._rp.cap_for(layer) + if self._stg is not None and self._spec > 0: + side_gen = self.read_gen() + # C++ demand_dual 是双源 decode 真实区的唯一权威（每层 1 次 inds 同步、零主线程落池/记账）。 + # 无 Python 退路：native 未编译 demand_dual → 明确报错（decode 依赖 native 全套能力）。 + if not getattr(self._rp, "_native_demand", False): + raise RuntimeError( + "双源 decode 需要 native demand_dual（已成唯一权威），但未检出。请编译 native_moe_ext。") + pool, local = self._acquire_native(layer, inds, side_gen, cap) + n_exp = cap + self._rp.spec_gens * self._rp.spec_slots + return pool, local, n_exp + pool, local = self._rp.acquire_gpu(layer, inds, num_experts) + return pool, local, cap + + def _native_meta(self, layer): + """缓存每层 demand_dual 的不变入参（pool_list/seg_nbytes/path），避免逐层重建 host 胶水。""" + cache = getattr(self, "_nd_meta", None) + if cache is None: + cache = self._nd_meta = {} + m = cache.get(layer) + if m is None: + stg = self._stg + segs = stg.src._segs # (proj, tensor, dt, shape, nb)，与池 key 同序 + pool_list = [self._rp._pools[layer][f"{p}.{t}"] for p, t, *_ in segs] + m = (pool_list, [int(nb) for *_, nb in segs], + f"{stg.src.dir}/layer{int(layer):02d}.blob", int(stg.stride)) + cache[layer] = m + return m + + def _acquire_native(self, layer, inds, side_gen, cap): + """方案B 取用：委派 C++ demand_dual（每层 1 次 inds 同步 + 并行 worker pread 落池）， + 更新 rp 统计计数（供报告口径一致）。""" + import mlx_streaming.native_moe_ext as N + rp = self._rp + rp._bootstrap_dual_pool(layer) # 首次建池 + real_init（幂等） + # 方案B 容量前提：miss 只能落真实区的 cap 槽，其中 pinned 永不可驱逐 → 可安放上限 cap−pinned。 + # 超了 C++ 会把专家落 0 号槽（拿别的专家权重算）→ 逐位不正确。调用方（block.py 的分流判据） + # 负责在超容量前把该前向送去 host/fetch；这里只做一次性告警兜底。 + _placeable = int(cap) - int(N.real_pinned_count(int(layer))) + if int(inds.size) > _placeable and not getattr(self, "_overcap_warned", False): + self._overcap_warned = True + import sys + print(f"[DEMAND_DUAL] 警告：inds.size={int(inds.size)} > 可安放槽 {_placeable}" + f"(cap={int(cap)} − pinned={int(cap) - _placeable})，真实区超容量，逐位将不正确。" + f"请调小 PREFILL_CHUNK，或调大 EXPERT_SLOTS / 调小 AUTOPIN_BUDGET_FRAC。", + file=sys.stderr, flush=True) + pool_list, seg_nbytes, path, stride = self._native_meta(layer) + local = N.demand_dual(inds, pool_list, seg_nbytes, int(layer), int(side_gen), path, + stride, int(cap), + rp.eviction_policy == "lfu", int(rp.lfu_decay_interval)) + st = N.demand_last_stats() # [hitpos, misspos, loads, fallback01, unplaced] + rp.hits += st[0] + rp.misses += st[2] + if len(st) > 4: + rp.unplaced += st[4] # >0 即有前向落了 0 号槽，输出不可信（/api/stats 暴露） + if st[3] == 0: + rp.gpu_fastpath += 1 + else: + rp.gpu_fallback += 1 + from mlx_streaming import config + if config.stg_verify(): # 诊断:方案B 池字节逐 key 真值校验(默认关) + self._verify_native_bytes(layer, inds, local) + return rp._pools[layer], local + + def _verify_native_bytes(self, layer, inds, local): + """诊断(STG_VERIFY，方案B)：校验「字节落池不变量」——真实区每个占用槽的池字节 == 该槽当前 + C++ 属主专家(g_real)的 blob 真值。这是 C++ 接管落池的字节等价铁证；发现不一致即池装错字节。 + + 注：不以 local→expert 为判据（local 可能因跨调用/多模型共享 g_real 而滞后于 g_real，属路由级 + 问题、非落池字节问题；逐位权威信号是 e2e n_mismatch）。 + """ + st = _stg_verify_state + st["calls"] += 1 + pool = self._rp._pools.get(layer) + if pool is None: + return + import mlx_streaming.native_moe_ext as N + stg = self._stg + path = f"{stg.src.dir}/layer{int(layer):02d}.blob" + segs = stg.src._segs + flat = N.real_region_contents(int(layer)) # [expert0,slot0,expert1,slot1,...] + for j in range(0, len(flat), 2): + e, slot = flat[j], flat[j + 1] + raw = N.blob_load(path, mx.array([e], dtype=mx.uint32), int(stg.stride))[0] + bad, off = None, 0 + for p, t, dt, shape, nb in segs: + k = f"{p}.{t}" + pv = pool[k][slot].reshape(-1).view(mx.uint8) + if not bool(mx.all(pv == raw[off:off + nb])): + bad = k + break + off += nb + if bad is None: + st["ok"] += 1 + else: + st["bad"] += 1 + if st["printed"] < 12: + st["printed"] += 1 + print(f"[STG_VERIFY-DUAL] BAD 落池字节错 call={st['calls']} layer={layer} " + f"expert={e} slot={slot} key={bad} (ok={st['ok']} bad={st['bad']})", flush=True) + + def acquire_host(self, layer, flat, inds_shape, inds_dtype, layer_cap): + """host/fetch 路径收口（prefill/大 seq 或关 GPU-remap）：flat 为 host 侧路由 id 列表。 + + 返回 (pool_arrays, local, n_experts)，与 block.py 原 host/fetch 分支逐元素等价： + - uniq <= cap：acquire(flat)，local 为槽位；n_experts = layer_cap。 + - uniq > cap：fetch(uniq_sorted)，local 为 remap 到 [0,uniq) 连续索引；n_experts = uniq 数。 + + native demand(dual)下真实区槽状态归 C++ g_real，Python 的 _slot_of/_free 从不填充， + 走 self._rp.acquire 必然在 _choose_victim 抛 "no evictable non-current slot"。 + 故 dual 下一律走 fetch 旁路（不进池、正确但慢）——这条路径本就是超容量前向的兜底。 + """ + import mlx.core as mx + cap = int(layer_cap) + uniq_set = set(flat) + # 入池的充分条件是 uniq ≤ cap − pinned:acquire 把本次全部唯一专家列为不可驱逐(current), + # 腾槽时只能挑「非 pinned 且非 current」的受害者,少一个都会在 _choose_victim 抛 + # "no evictable non-current slot"。pinned 是那层实际钉死数,不是拍出来的固定余量—— + # 写死余量会在 cap 小于余量时把所有请求推进 fetch 慢路径(每次重读盘)。 + _pinned_n = len(getattr(self._rp, "_pinned", {}).get(layer, ()) or ()) + if (not getattr(self._rp, "_native_demand", False) + and len(uniq_set) <= max(1, cap - _pinned_n)): + pool, slots = self._rp.acquire(layer, flat) + local = mx.array(slots, dtype=inds_dtype).reshape(inds_shape) + return pool, local, cap + uniq_sorted = sorted(uniq_set) + remap = {g: i for i, g in enumerate(uniq_sorted)} + local = mx.array([remap[i] for i in flat], dtype=inds_dtype).reshape(inds_shape) + # 超容量(或 dual:真实区槽状态归 C++ g_real,Python _slot_of 从不填充 → acquire 必抛) + # 走 fetch 旁路:不进池、按 uniq 现堆一份参与计算,正确但慢。这条路径本就是兜底。 + # fetch 由 blob 感知的 FileExpertStore 提供;ResidentExpertPool 本身没有,故优先用注入的 store。 + src = self._store if self._store is not None else self._rp + if not hasattr(src, "fetch"): + raise RuntimeError("VirtualPool.acquire_host 超容量回退需要带 fetch 的 store(构造时传入)") + fetched = src.fetch(layer, uniq_sorted) + return fetched, local, len(uniq_sorted) + + def prefetch(self, layer, pred, resident, pool_list): + """向 fill 代 submit 预读：base_row = cap_for(layer) + fill_gen*spec。""" + g = self.fill_gen() + base = self._rp.cap_for(layer) + g * self._spec + return self._stg.submit_pool_sideregion(layer, pred, resident, pool_list, base, gen=g) diff --git a/mlx_streaming/core/linear_attn/__init__.py b/mlx_streaming/core/linear_attn/__init__.py new file mode 100644 index 0000000..ddda1f0 --- /dev/null +++ b/mlx_streaming/core/linear_attn/__init__.py @@ -0,0 +1 @@ +"""线性注意力（Qwen3-Next gated-delta）相关的自定义 kernel。""" diff --git a/mlx_streaming/core/linear_attn/gated_delta_multistate.py b/mlx_streaming/core/linear_attn/gated_delta_multistate.py new file mode 100644 index 0000000..634132c --- /dev/null +++ b/mlx_streaming/core/linear_attn/gated_delta_multistate.py @@ -0,0 +1,226 @@ +"""逐步状态输出版 gated-delta Metal kernel（vendored 自 mlx_lm.models.gated_delta）。 + +来源：mlx-lm 0.x `mlx_lm/models/gated_delta.py` 的 `_make_gated_delta_kernel` +（本机 mlx 0.31.2）。本文件**只**在mlx kernel 基础上多加一个输出 `states_out`： +在时间维 `for (t=0..T-1)` 的串行递归里，每算完一步就把当前 state 写回 +`states_out[B, T, Hv, Dv, Dk]`。 + +为什么 bit-exact：mlx kernel 的时间维本就是线程内串行递归（不是分块并行），第 t 步 +寄存器里的 `state` 就是「逐 token 走 t+1 步后的状态」，全程 fp32、同序、同融合。因此 +一次 seq=T 调用产出的 `states_out[:, t]`，与逐 token（T 次 seq=1）调用链得到的 state +逐 bit 相等。这正是 MTP 投机验证后能「按接受长度直接提交、零 replay」的数值根基。 + +用途：MTP 验证前向里捕获每个 token 后的 ssm 递归态作为 checkpoint，替换原先走 +`_gated_delta_step_ops`（与 baseline kernel 不 bit-exact）的捕获路径。 + +非目标：仅 Metal(GPU) 路径；CPU/ops 回退不在此实现（部署恒为 Metal）。 +""" +from typing import Optional, Tuple + +import mlx.core as mx +import mlx.nn as nn + +# 与 mlx 一致：g = exp(-exp(A_log) * softplus(a + dt_bias))，fp32。 +from mlx_lm.models.gated_delta import compute_g + + +def _make_gated_delta_multistate_kernel(has_mask=False, vectorized=False): + """构造逐步状态输出版 kernel；除多写 states_out 外与mlx kernel 完全一致。""" + if not mx.metal.is_available(): + return None + mask_source = "mask[b_idx * T + t]" if has_mask else "true" + + # g 索引方式：标量门控 [B,T,Hv] vs 向量门控 [B,T,Hv,Dk]，与 mlx 对齐。 + if vectorized: + g_comment = "// g: [B, T, Hv, Dk]" + g_setup = "auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;" + g_access = "g_[s_idx]" + g_advance = "g_ += Hv * Dk;" + else: + g_comment = "// g: [B, T, Hv]" + g_setup = "auto g_ = g + b_idx * T * Hv;" + g_access = "g_[hv_idx]" + g_advance = "g_ += Hv;" + + # 与 mlx 版唯一的差异：时间循环每步末把 state 写入 states_out（见下方标注的新增块）。 + source = f""" + 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; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + 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; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (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]); + }} + + {g_comment} + {g_setup} + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) {{ + if ({mask_source}) {{ + 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_access}; + 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 (thread_index_in_simdgroup == 0) {{ + y[dv_idx] = static_cast(out); + }} + }} else {{ + y[dv_idx] = static_cast(0); + }} + + // ===== 新增：每步把当前 state 写回 states_out[B, T, Hv, Dv, Dk] ===== + // masked 步 state 不变（上面 else 分支不更新 state），写回的就是上一步的状态， + // 与逐 token 单步链语义一致。 + auto ms_ = states_out + + ((((b_idx * T + t) * Hv + hv_idx) * Dv + dv_idx) * Dk); + for (int i = 0; i < n_per_t; ++i) {{ + auto s_idx = n_per_t * dk_idx + i; + ms_[s_idx] = static_cast(state[i]); + }} + // =================================================================== + + // Increment data pointers to next time step + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + {g_advance} + beta_ += Hv; + }} + for (int i = 0; i < n_per_t; ++i) {{ + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + }} + """ + inputs = ["q", "k", "v", "g", "beta", "state_in", "T"] + if has_mask: + inputs.append("mask") + + suffix = "" + if vectorized: + suffix += "_vec" + if has_mask: + suffix += "_mask" + + return mx.fast.metal_kernel( + name=f"gated_delta_step_multistate{suffix}", + input_names=inputs, + output_names=["y", "state_out", "states_out"], + source=source, + ) + + +_kernel = _make_gated_delta_multistate_kernel(has_mask=False, vectorized=False) +_kernel_masked = _make_gated_delta_multistate_kernel(has_mask=True, vectorized=False) +_kernel_vec = _make_gated_delta_multistate_kernel(has_mask=False, vectorized=True) +_kernel_vec_masked = _make_gated_delta_multistate_kernel(has_mask=True, vectorized=True) + + +def gated_delta_multistate_kernel( + q: mx.array, + k: mx.array, + v: mx.array, + g: mx.array, + beta: mx.array, + state: mx.array, + mask: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array, mx.array]: + """逐步状态输出版 kernel 调用，返回 (y, final_state, states_out)。 + + 形状（与 mlx `gated_delta_kernel` 一致，不在 Python 侧做 GQA repeat）： + - q, k: [B, T, Hk, Dk] + - v: [B, T, Hv, Dv] + - g: [B, T, Hv]（标量）或 [B, T, Hv, Dk]（向量） + - beta: [B, T, Hv] + - state:[B, Hv, Dv, Dk] + 返回： + - y: [B, T, Hv, Dv] + - final_state:[B, Hv, Dv, Dk]（== states_out[:, -1]） + - states_out: [B, T, Hv, Dv, Dk]（每个 token 处理完后的递归态） + """ + B, T, Hk, Dk = k.shape + Hv, Dv = v.shape[2:] + input_type = q.dtype + state_type = state.dtype + if g.ndim == 4: + kernel = _kernel_vec_masked if mask is not None else _kernel_vec + else: + kernel = _kernel_masked if mask is not None else _kernel + + inputs = [q, k, v, g, beta, state, T] + if mask is not None: + inputs.append(mask) + + return kernel( + inputs=inputs, + template=[ + ("InT", input_type), + ("StT", state_type), + ("Dk", Dk), + ("Dv", Dv), + ("Hk", Hk), + ("Hv", Hv), + ], + grid=(32, Dv, B * Hv), + threadgroup=(32, 4, 1), + output_shapes=[(B, T, Hv, Dv), state.shape, (B, T, Hv, Dv, Dk)], + output_dtypes=[input_type, state_type, state_type], + ) + + +def gated_delta_update_multistate( + q: mx.array, + k: mx.array, + v: mx.array, + a: mx.array, + b: mx.array, + A_log: mx.array, + dt_bias: mx.array, + state: Optional[mx.array] = None, + mask: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array, mx.array]: + """对齐 mlx `gated_delta_update` 的签名，多返回逐步状态 states_out。 + + 内部 beta/g 计算与 mlx 完全一致（sigmoid(b) / compute_g），保证数值路径相同。 + 仅走 Metal kernel 路径（multistate 仅 GPU 实现）。 + """ + beta = mx.sigmoid(b) + g = compute_g(A_log, a, dt_bias) + if state is None: + B, _, Hk, Dk = q.shape + Hv, Dv = v.shape[-2:] + state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32) + return gated_delta_multistate_kernel(q, k, v, g, beta, state, mask) diff --git a/mlx_streaming/core/mem.py b/mlx_streaming/core/mem.py new file mode 100644 index 0000000..4d8eec3 --- /dev/null +++ b/mlx_streaming/core/mem.py @@ -0,0 +1,75 @@ +"""内存度量统一口径。macOS 上 ru_maxrss 单位是字节。""" +import resource +from dataclasses import dataclass + +import mlx.core as mx + + +def rss_bytes() -> int: + # macOS: ru_maxrss 已是字节；Linux 是 KB（本项目目标是 macOS） + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + +def _call(name: str) -> int: + fn = getattr(mx, name, None) + if fn is None: + fn = getattr(getattr(mx, "metal", object()), name, None) + try: + return int(fn()) if fn else 0 + except Exception: + return 0 + + +@dataclass +class MemSnapshot: + rss_bytes: int + mlx_active_bytes: int + mlx_peak_bytes: int + + +def snapshot() -> MemSnapshot: + return MemSnapshot( + rss_bytes=rss_bytes(), + mlx_active_bytes=_call("get_active_memory"), + mlx_peak_bytes=_call("get_peak_memory"), + ) + + +def clear_cache() -> None: + fn = getattr(mx, "clear_cache", None) or getattr(getattr(mx, "metal", object()), "clear_cache", None) + if fn: + fn() + + +def reset_peak() -> None: + fn = getattr(mx, "reset_peak_memory", None) or getattr(getattr(mx, "metal", object()), "reset_peak_memory", None) + if fn: + fn() + + +def _set_limit(name: str, nbytes: int) -> bool: + """调用 mx.set_*_limit(字节);存在且成功返回 True。""" + fn = getattr(mx, name, None) or getattr(getattr(mx, "metal", object()), name, None) + if fn is None: + return False + try: + fn(int(nbytes)) + return True + except Exception: + return False + + +def setup_memory_hygiene(cache_gb: float = 2.0, wired_gb: float = 0.0) -> dict: + """长期运行的内存防御:bound MLX 缓冲缓存 + 可选 wire 工作集防 macOS 压缩器。 + + - cache_gb>0:set_cache_limit,封顶 MLX 可回收缓冲,防长会话里缓存膨胀把常驻推过墙。 + - wired_gb>0:set_wired_limit,把这么多 GB 的 GPU 缓冲钉为常驻(wired),macOS 不再 + 压缩/换出这些页 → 长跑延迟稳定。务必 < 系统建议工作集(本机 26.8GB),否则饿死系统。 + 返回实际生效项,便于启动日志记录。 + """ + applied = {} + if cache_gb and cache_gb > 0: + applied["cache_limit_gb"] = cache_gb if _set_limit("set_cache_limit", cache_gb * 1e9) else None + if wired_gb and wired_gb > 0: + applied["wired_limit_gb"] = wired_gb if _set_limit("set_wired_limit", wired_gb * 1e9) else None + return applied diff --git a/mlx_streaming/core/moe/__init__.py b/mlx_streaming/core/moe/__init__.py new file mode 100644 index 0000000..cdbd4e3 --- /dev/null +++ b/mlx_streaming/core/moe/__init__.py @@ -0,0 +1 @@ +"""MoE 模块：门控选专家(gate) / 专家计算(compute) / 自定义算子(custom_kernel) / 热路径块(block)。""" diff --git a/mlx_streaming/core/moe/block.py b/mlx_streaming/core/moe/block.py new file mode 100644 index 0000000..bf6850c --- /dev/null +++ b/mlx_streaming/core/moe/block.py @@ -0,0 +1,431 @@ +"""MoE 热路径块：路由 → 选专家 → 取专家权重 → 计算 → 加权合并。 + +包含两种块： +- `StreamingMoeBlock`：包住原生 MoE 块，专家权重常驻（switch_mlp），只激活选中专家。 +- `FileStreamingMoeBlock`：专家权重从磁盘按需加载（流式），是低内存推理的核心热路径， + 集成常驻池 acquire、native-fused-prefetch 预取、共享专家叠加等逻辑。 +""" +import os +import time + +import mlx.core as mx + +from mlx_streaming import config +from mlx_streaming.core.moe import native_moe +from mlx_streaming.core import route_trace +from mlx_streaming.core.profiling import ( + PROF, WINDOW_PROF, PREDICT_RECALL_PROF, MISS_ATTRIB, note_miss_attrib, + note_tprof, TPROF_ON, _PROF_ON, _tick, note_union, UNION_ON) + +# decode/verify 热路径判据:seq 短(单 token decode=1、MTP verify=K≤几)，与 prefill 长 seq 区分。 +_DECODE_SEQ_MAX = 8 +from mlx_streaming.core.moe.gate import _effective_top_k +from mlx_streaming.core.cache import autopin +from mlx_streaming.core.moe.compute import ( + streaming_switch_glu_forward, PersistentSubGLU) + + +class StreamingMoeBlock: + """包住原 Qwen3MoeSparseMoeBlock：路由器常驻，专家计算改为只算选中专家。 + + 若提供 store（LruExpertStore），可进一步把专家权重从磁盘按需加载；否则直接在 + 常驻的 switch_mlp 上做 uniq 切片计算（仍只激活少数专家）。 + """ + + def __init__(self, orig_block, layer_idx: int, store=None): + self.gate = orig_block.gate # 路由器常驻（很小） + self.top_k = orig_block.top_k + self.norm_topk_prob = orig_block.norm_topk_prob + self.switch_mlp = orig_block.switch_mlp + self.store = store + self.layer_idx = layer_idx + + def __call__(self, x: mx.array) -> mx.array: + gates = mx.softmax(self.gate(x), axis=-1, precise=True) + k = _effective_top_k(self.top_k) + inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:] + scores = mx.take_along_axis(gates, inds, axis=-1) + if self.norm_topk_prob: + scores = scores / mx.sum(scores, axis=-1, keepdims=True) + mx.eval(inds) # 先物化路由结果，才能按结果取专家 + y = streaming_switch_glu_forward(self.switch_mlp, x, inds) + return (y * scores[..., None]).sum(axis=-2) + + +class FileStreamingMoeBlock: + """文件后端流式 MoE 块：路由器常驻，专家权重从磁盘按需加载（不持有堆叠 switch_mlp）。""" + + def __init__(self, gate, top_k, norm_topk_prob, store, layer_idx, + hidden, moe_inter, group_size, bits, + proj_bits: dict | None = None, + shared_expert=None, shared_expert_gate=None): + self.gate = gate + self.top_k = top_k + self.norm_topk_prob = norm_topk_prob + self.store = store + self.layer_idx = layer_idx + self.hidden = hidden + self.moe_inter = moe_inter + self.group_size = group_size + self.bits = bits + # Qwen3-Next 等带共享专家的模型：共享专家恒激活、必须常驻（不流式）， + # 输出叠加 sigmoid(shared_expert_gate(x)) * shared_expert(x)。None 则退化为纯路由（如 Qwen3-MoE）。 + self.shared_expert = shared_expert + self.shared_expert_gate = shared_expert_gate + # 可选全流式 blob 源（STREAM_BLOB=1 时由 model_builder 注入），见 core/blob_loader.py。 + self._blob = None + # 持久化子模块：跨 token 复用，避免每调用重建 QSL。 + # proj_bits 非空时走混合精度（逐 proj 不同 bit）。 + self._sub = PersistentSubGLU(hidden, moe_inter, group_size, bits, + proj_bits=proj_bits, layer_idx=layer_idx) + + def __call__(self, x: mx.array) -> mx.array: + if _PROF_ON: + return self._call_prof(x) + gates = mx.softmax(self.gate(x), axis=-1, precise=True) + k = _effective_top_k(self.top_k) + inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:] + scores = mx.take_along_axis(gates, inds, axis=-1) + if self.norm_topk_prob: + scores = scores / mx.sum(scores, axis=-1, keepdims=True) + if config.predict_recall_prof(): + ps = getattr(self, "_predicted_set", None) + if ps is not None: + act = {int(i) for i in inds.reshape(-1).tolist()} + PREDICT_RECALL_PROF["hit"] += len(ps & act) + PREDICT_RECALL_PROF["routed"] += len(act) + PREDICT_RECALL_PROF["n"] += 1 + if config.native_moe(): + native_y = self._try_native_forward(x, inds, scores, gates.shape[-1]) + if native_y is not None: + y = native_y + if self.shared_expert is not None: + y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x) + return y + if self._blob is not None and config.stream_blob(): + # 全流式 blob 路径：每层按需并行读专家 → 复用 _sub.forward(MLX quantized_matmul)。 + flat = [int(i) for i in inds.reshape(-1).tolist()] + pool_arrays, slots = self._blob.acquire(self.layer_idx, flat) + local = mx.array(slots, dtype=inds.dtype).reshape(inds.shape) + y = self._sub.forward(pool_arrays, len(set(flat)), x, local) + y = (y * scores[..., None]).sum(axis=-2) + if self.shared_expert is not None: + y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x) + return y + if config.probe_perlayer_sync(): + # 诊断：每层强制一次 host 同步（模拟预测 barrier），不做任何预取。 + # 若单这个就拖慢 ~30% → 税是 eval barrier，C++ 读统一内存也救不了。 + _ = int(inds.reshape(-1)[:1].item()) + if (config.stream_blob_bg() + and getattr(self.store, "_bg", None) is not None): + # 后台预取已物化的专家在此（主线程、acquire 前）写进常驻池槽 → 转 miss 为 hit。 + if config.window_prof(): + t0 = getattr(self, "_submit_t", None) + if t0 is not None: + WINDOW_PROF["sum_s"] += time.perf_counter() - t0 + WINDOW_PROF["n"] += 1 + self.store.promote_prefetched(self.layer_idx) + # native-fused-prefetch 的 promote（把回调预读好的专家写进池槽）下移到各 acquire 分支前： + # host/verify 路径已算出真实路由 uniq_set，可"只 promote 命中本层路由的专家"，零额外同步地 + # 丢弃假阳性 → 省掉无用 scatter + 池污染（这是 promote -3.2% 开销的主因）。 + _stg_mgr = getattr(self.store, "_staging", None) + # 双源：投机专家留池侧区由 demand_dual 取，不 promote、不驱逐 → 关 promote。 + _do_promote = (_stg_mgr is not None and not config.native_no_promote() + and not config.zerocopy_dual_source()) + if config.route_trace_enabled(): + flat_trace = [int(i) for i in inds.reshape(-1).tolist()] + resident = (self.store.resident_experts(self.layer_idx) + if hasattr(self.store, "resident_experts") else set()) + rank = (self.store.resident_lru_scores(self.layer_idx) + if hasattr(self.store, "resident_lru_scores") else {}) + route_trace.record( + self.layer_idx, flat_trace, set(flat_trace) - resident, resident, rank) + # native-fused-prefetch：在 seq 分支**之前**提交，使 decode(seq=1) 与 MTP verify(seq=K) + # 两条路径都触发预取。dummy 折进 inds(加 0)：GPU 路径靠 acquire_gpu 的 n_miss eval、 + # host 路径靠后面的 .tolist() eval，都会触发完成回调里的 pread。verify 时 x 为 K 个 + # token（verify_in=[x, d_1..d_{K-1}]），预测的是"下一层这 K 个 token 的专家并集"（recall≈0.96 的口径）。 + # 双源双缓冲：每块开头先推进/检测前向边界（必须在本块预取提交之前，保证同前向内 fill_gen/read_gen 恒定）。 + if config.zerocopy_dual_source() and getattr(self, "_vpool", None) is not None: + self._vpool.begin_forward(self.layer_idx) + if (getattr(self.store, "_staging", None) is not None + and not config.native_no_submit()): + _dummy = self._native_fused_prefetch(x) + if _dummy is not None: + inds = inds + (_dummy.reshape(()).astype(inds.dtype) * 0) + layer_cap = self.store.cap_for(self.layer_idx) + # verify(小 seq)可走 GPU 重映射;prefill(大 seq)唯一专家可能超 cap,须留在 host 路径(有超容量 fetch 回退)。 + if config.zerocopy_dual_source() and getattr(self, "_vpool", None) is not None: + # 双源两级缓存本就是为 MTP verify 建的:verify 走 dual 路径才能读到侧区,否则落 host + # 路径会白填侧区(预取填了却不读)→ 命中骤降、读盘翻倍。 + # + # 但判据必须按 demand_dual 真实区**可安放容量**算,不能按 cap+侧区行:侧区行是预取只读的, + # miss 只能落真实区的 cap 槽,其中 pinned(AUTOPIN)永不可驱逐。设本次唯一专家数 U、 + # 真实区 cap=C、pinned=P,可用槽 = C − P − (本次命中的非 pinned 常驻数),推得 + # 「U ≤ C − P」是安放得下的充分条件。超了 C++ 分不到槽会把专家落 0 号槽 → 逐位错算 + # (prefill chunk 大时曾整段 system+tools 都算错)。U 无同步上界即 inds.size = seq×top_k。 + _verify_gpu = (x.shape[1] * k <= self._dual_placeable(layer_cap)) + else: + _verify_gpu = (config.verify_gpu_remap() and x.shape[1] * k <= layer_cap) + if (config.resident_pool_enabled() + and config.gpu_remap_enabled() + and (x.shape[1] == 1 or _verify_gpu)): + # decode 热路径:GPU 侧 slot 重映射,命中层零 host 往返(消除每层 .tolist 同步), + # 仅一次 miss 标志同步;真 miss 才回退读盘。decode top-k(≤cap)恒装得下。 + # 专家总数 = gate 输出维(gates 末维),无需依赖 gate.weight。 + if UNION_ON: + # GPU 路径本不 .tolist 真实路由;仅 UNION_PROF=1 时付一次同步取并集大小。 + note_union(x.shape[1], len({int(i) for i in inds.reshape(-1).tolist()}), + self.layer_idx) + # decode/verify GPU 重映射路径:host 无现成真实路由。开关开时用 GPU membership 仅对 + # 预读候选现算 used(drain≤budget)过滤假阳性;关时退回整批写入(回退/对照基线)。 + if _do_promote: + if config.gpu_remap_promote_filter(): + _stg_mgr.promote(self.layer_idx, self.store, + route_inds=inds, num_experts=gates.shape[-1]) + else: + _stg_mgr.promote(self.layer_idx, self.store) + if config.miss_attrib(): + # 诊断专用:GPU remap 路径本不 .tolist 真实路由(那是它消栅栏的关键), + # 仅 MISS_ATTRIB=1 时付一次 .tolist 取真实路由,统计 decode 热路径 A/B 构成。 + _uniq = {int(i) for i in inds.reshape(-1).tolist()} + _ps = getattr(self, "_predicted_set", None) + _res = (self.store.resident_experts(self.layer_idx) + if hasattr(self.store, "resident_experts") else set()) + _rdy = _stg_mgr.last_ready.get(self.layer_idx) if _stg_mgr else None + note_miss_attrib(_uniq, _ps, _res, x.shape[1] <= _DECODE_SEQ_MAX, _rdy) + if config.zerocopy_dual_source(): + # 双源双缓冲：统一走 VirtualPool.acquire（呈现「所有专家都在」视角， + # 内部真实区表 ∪ 侧区(读代) 单次 gather，返回 (pool, local, n_experts)）。 + pool_arrays, local, n_experts = self._vpool.acquire( + self.layer_idx, inds, gates.shape[-1], + seq_len=x.shape[1], layer_cap=layer_cap) + y = self._sub.forward(pool_arrays, n_experts, x, local) + else: + pool_arrays, local = self.store.acquire_gpu( + self.layer_idx, inds, gates.shape[-1]) + if config.stg_verify(): + # 诊断：消费侧字节校验（默认 off），定位混合 ahead 下池槽损坏。 + self.store._resident.verify_acquire_bytes( + self.layer_idx, inds, _stg_mgr) + y = self._sub.forward(pool_arrays, layer_cap, x, local) + else: + # prefill/大批量(seq>1)或显式关闭 GPU_REMAP:走 host 路径。 + # 只在此处对 inds 做一次 .tolist() 同步，uniq/local 全在 Python 里算。 + flat = [int(i) for i in inds.reshape(-1).tolist()] + autopin.note(self.layer_idx, flat) # AUTOPIN 热度:搭已物化的 flat,零额外同步(AUTOPIN=0 立即返回) + uniq_set = set(flat) + if UNION_ON: + note_union(x.shape[1], len(uniq_set), self.layer_idx) # 本层路由专家并集(零额外同步,uniq_set 已算) + # promote 只写"预读好 ∩ 本层真实路由"的专家：复用已算的 uniq_set，零额外同步， + # 假阳性不进池 → 省 scatter、不污染池（acquire 前完成，命中转化生效）。 + if _do_promote: + _stg_mgr.promote(self.layer_idx, self.store, used=uniq_set) + if config.miss_attrib(): + # 在 promote 之后、acquire(读盘)之前统计：本层真实路由的命中构成。 + # miss 拆成 A(预测到却没进池：budget丢/时序/驱逐) 与 B(没预测到：召回缺口)。 + _ps = getattr(self, "_predicted_set", None) + _res = (self.store.resident_experts(self.layer_idx) + if hasattr(self.store, "resident_experts") else set()) + _rdy = _stg_mgr.last_ready.get(self.layer_idx) if _stg_mgr else None + note_miss_attrib(uniq_set, _ps, _res, x.shape[1] <= _DECODE_SEQ_MAX, _rdy) + # 池只能同时容纳 ≤该层容量 个唯一专家；prefill 唯一数超容量时回退 stack。 + # dual 模式（_vpool 存在）统一走 VirtualPool.acquire_host 收口 host+fetch； + # 非 dual（_vpool 为 None）保留内联逻辑（非目标路径，避免额外依赖）。 + if (config.resident_pool_enabled() + and getattr(self, "_vpool", None) is not None): + pool_arrays, local, n_experts = self._vpool.acquire_host( + self.layer_idx, flat, inds.shape, inds.dtype, layer_cap) + y = self._sub.forward(pool_arrays, n_experts, x, local) + elif (config.resident_pool_enabled() + and len(uniq_set) <= layer_cap): + pool_arrays, slots = self.store.acquire(self.layer_idx, flat) + local = mx.array(slots, dtype=inds.dtype).reshape(inds.shape) + y = self._sub.forward(pool_arrays, layer_cap, x, local) + else: + uniq_sorted = sorted(uniq_set) + remap = {g: i for i, g in enumerate(uniq_sorted)} + local = mx.array([remap[i] for i in flat], dtype=inds.dtype).reshape(inds.shape) + fetched = self.store.fetch(self.layer_idx, uniq_sorted) + y = self._sub.forward(fetched, len(uniq_sorted), x, local) + y = (y * scores[..., None]).sum(axis=-2) + if self.shared_expert is not None: # 共享专家（常驻）叠加 + y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x) + return y + + def _dual_placeable(self, layer_cap: int) -> int: + """demand_dual 真实区本层能安放的唯一专家上限 = cap − pinned。 + + pinned 集合只在启动 AUTOPIN 预热时写入、生成期恒定,故首次查到后缓存 + (decode 每层每步都要判,不能每次都进 C++ 取锁)。 + """ + n = getattr(self, "_pinned_n", None) + if n is None: + try: + import mlx_streaming.native_moe_ext as _N + n = int(_N.real_pinned_count(int(self.layer_idx))) + except Exception: # noqa: BLE001 取不到就按最保守的「全钉死」算,逼回 host 路径 + n = int(layer_cap) + self._pinned_n = n + return max(1, int(layer_cap) - n) + + def _native_fused_prefetch(self, x: mx.array): + """搭车式预取：用下 AHEAD 层 gate 对 x 算预测 inds(lazy)，挂 GPU 完成回调， + 在 C++ 里(GPU 算完后)读 id + pread 预热下层专家字节。返回 dummy 张量(需被 eval 才触发)。 + 全程不加主线程 host 同步——靠折进 inds、搭 acquire_gpu 的 n_miss eval。 + """ + model = getattr(self, "_prefetch_model_ref", None) + bl = getattr(self.store, "_blob_loader", None) + if model is None or bl is None: + return None + layers = getattr(model, "layers", []) + vp = getattr(self, "_vpool", None) + if vp is not None: + tgt = vp.target_for(self.layer_idx) # per-layer cutoff ahead + if not (self.layer_idx < tgt < len(layers)): # 0/越界/无前瞻 → 跳过 + return None + else: + ahead = max(1, config.cross_layer_ahead(default=1)) + tgt = self.layer_idx + ahead + if not (0 <= tgt < len(layers)): + return None + tmlp = getattr(layers[tgt], "mlp", None) + if not isinstance(tmlp, FileStreamingMoeBlock): + return None + try: + from mlx_streaming import native_moe_ext as _N + except Exception: + return None + # 方案B：预测"宽集合"(top-N, N=predict_width)只为高 recall，不占内存（仅一次 gate argpartition）； + # 真正占 staging 的是 C++ 回调过滤常驻后、按分截到 staging budget 的"缺口"子集。 + _tp0 = time.perf_counter() if TPROF_ON else 0.0 + predict_width = config.cross_layer_predict_width() + # 默认(PREDICT_USE_X=1)：用本层 MoE 输入 x 喂目标层 gate。x 含了本层 attention，离 L+1 + # 近半层 → 比未归一化输入更新鲜，实测 recall 0.812→0.847(+3.6pp)、还省一次 norm。 + # PREDICT_USE_X=0 回退旧路径：目标层 post_attention_layernorm 作用在本层未归一化输入上。 + h = getattr(self, "_unnormed_input", None) + if config.predict_use_x(): + g = tmlp.gate(x) + elif h is None: + h = x # 退化：无未归一化输入时用 x（norm 偏差，覆盖会差） + g = tmlp.gate(h) + else: + g = tmlp.gate(layers[tgt].post_attention_layernorm(h)) + # 按 seq 维聚合成"该批 K 个 token 的专家并集"近似。聚合方式 PREDICT_AGG： + # - max（默认）：任一 token 强烈想要即入选；mean：各 token 平均偏好； + # - union：每 token 各取 top-kk 再并集（与真实路由"并集"结构最一致，候选更多，由 handler 去重+截断）。 + _agg = config.predict_agg() + kk = min(g.shape[-1], predict_width) + if g.ndim == 3 and _agg == "union": + # union 候选数 = seq × union_k。每 token 只取 top-union_k（独立于 predict_width）， + # 调小 union_k 缩候选 → 减少缺口、缓解 staging/读盘洪水（每 token 真实只路由 top_k 个）。 + uk = min(g.shape[-1], config.predict_union_k()) + pred = mx.argpartition(g, kth=-uk, axis=-1)[..., -uk:].reshape(-1).astype(mx.uint32) + else: + if g.ndim == 3: + g = g.mean(axis=1) if _agg == "mean" else g.max(axis=1) + # 取 top-kk（argpartition，O(E) 不全排序）。实测：加宽 kk 不提升命中——staging 更多 = 后台 + # pread 更多、到位更晚 + 抢带宽 → 反而更低；故 width=16 即最优，cap 截断极少触发、无需排序。 + pred = mx.argpartition(g, kth=-kk, axis=-1)[..., -kk:].reshape(-1).astype(mx.uint32) + if config.predict_recall_prof() or config.miss_attrib(): + try: + tmlp._predicted_set = {int(e) for e in pred.tolist()} + except Exception: + pass + if TPROF_ON: + note_tprof("predict_s", time.perf_counter() - _tp0, count_key="predict_n") + stg = getattr(self.store, "_staging", None) + if stg is not None: + resident = (self.store.resident_experts(tgt) + if hasattr(self.store, "resident_experts") else None) + if config.zerocopy_dual_source(): + # 双源双缓冲：经 VirtualPool 向填代 submit,C++ 回调把预读段散写进该代侧区行。 + rp = self.store._resident + if tgt not in rp._pools: + return None # 目标层池未建(首 token 预热) → 跳过本次 + segs = stg.src._segs # (proj, tensor, dt, shape, nb)，与池 key 顺序一致 + pool_list = [rp._pools[tgt][f"{p}.{t}"] for p, t, *_ in segs] + import os as _os + if _os.environ.get("POOL_PTR_TRACE") is not None and tgt == int(_os.environ["POOL_PTR_TRACE"]): + try: + import mlx_streaming.native_moe_ext as _N + _fk = f"{segs[0][0]}.{segs[0][1]}" + print(f"[POOL_PTR] submit layer={tgt} key={_fk} obj_id={id(rp._pools[tgt][_fk])} " + f"ptr={hex(_N.array_data_ptr(rp._pools[tgt][_fk]))}", flush=True) + except Exception as _e: + print(f"[POOL_PTR] err {_e}", flush=True) + return self._vpool.prefetch(tgt, pred, resident, pool_list) + # miss→hit:回调按目标层常驻快照过滤，只把缺口 pread 进 staging（≤budget 行），promote 时写池。 + if TPROF_ON: + _ts0 = time.perf_counter() + _r = stg.submit(tgt, pred, resident) + note_tprof("submit_s", time.perf_counter() - _ts0, count_key="submit_n") + return _r + return stg.submit(tgt, pred, resident) + # 仅预热字节（page cache）的轻量版。 + path = os.path.join(bl.dir, f"layer{tgt:02d}.blob") + return _N.prefetch_on_complete(pred, path, int(bl.stride), True) + + def _try_native_forward(self, x: mx.array, inds: mx.array, scores: mx.array, num_experts: int): + # native 第一版只支持三投影同 bit 的专家；其他情况保持现有 MLX 回退。 + proj_bits = getattr(self._sub, "proj_bits", {}) + bits_set = {int(proj_bits.get(name, self.bits)) for name in ("gate_proj", "up_proj", "down_proj")} + if len(bits_set) != 1: + return None + bits = bits_set.pop() + if not native_moe.can_native_moe( + self.layer_idx, self.hidden, self.moe_inter, + self.group_size, bits, num_experts): + return None + t_sync = time.perf_counter() + mx.eval(inds) + flat = [int(i) for i in inds.reshape(-1).tolist()] + native_moe.note_route_sync(time.perf_counter() - t_sync) + return native_moe.try_native_moe( + self.layer_idx, flat, x, scores, + self.hidden, self.moe_inter, self.group_size, bits, num_experts) + + def _call_prof(self, x: mx.array) -> mx.array: + t = time.perf_counter() + gates = mx.softmax(self.gate(x), axis=-1, precise=True) + k = _effective_top_k(self.top_k) + inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:] + scores = mx.take_along_axis(gates, inds, axis=-1) + if self.norm_topk_prob: + scores = scores / mx.sum(scores, axis=-1, keepdims=True) + mx.eval(inds, scores) + _tick("route", t); t = time.perf_counter() + + flat = [int(i) for i in inds.reshape(-1).tolist()] + uniq_set = set(flat) + layer_cap = self.store.cap_for(self.layer_idx) + use_pool = (config.resident_pool_enabled() + and len(uniq_set) <= layer_cap) + if not use_pool: + uniq_sorted = sorted(uniq_set) + remap = {g: i for i, g in enumerate(uniq_sorted)} + local = mx.array([remap[i] for i in flat], dtype=inds.dtype).reshape(inds.shape) + mx.eval(local) + _tick("pyremap", t); t = time.perf_counter() + + if use_pool: + pool_arrays, slots = self.store.acquire(self.layer_idx, flat) + local = mx.array(slots, dtype=inds.dtype).reshape(inds.shape) + mx.eval(local) + fetched, n_experts = pool_arrays, layer_cap + else: + fetched = self.store.fetch(self.layer_idx, uniq_sorted) + n_experts = len(uniq_sorted) + mx.eval(list(fetched.values())) + _tick("fetch", t); t = time.perf_counter() + + y = self._sub.forward(fetched, n_experts, x, local) + mx.eval(y) + _tick("matmul", t); t = time.perf_counter() + + y = (y * scores[..., None]).sum(axis=-2) + if self.shared_expert is not None: + y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x) + mx.eval(y) + _tick("combine", t) + PROF["n_calls"] += 1 + return y diff --git a/mlx_streaming/core/moe/compute.py b/mlx_streaming/core/moe/compute.py new file mode 100644 index 0000000..80bd4ec --- /dev/null +++ b/mlx_streaming/core/moe/compute.py @@ -0,0 +1,208 @@ +"""MoE 专家计算：把选中专家切片成小 SwitchGLU 并前向（数值等价全量 SwitchGLU）。 + +设计要点： +- `_unique_and_local`：把 top-k 路由的全局专家 id 压成「唯一集合 + 本地下标」， + 只在实际激活的少数专家上 gather 计算。 +- `PersistentSubGLU`：跨调用复用一套 SwitchGLU + 3×QuantizedSwitchLinear 对象， + 解码 batch=1 时唯一专家数恒为 top_k，只首次构造、之后原地 update 权重，省重建开销。 +""" +from typing import Tuple + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.switch_layers import SwitchLinear, QuantizedSwitchLinear, SwitchGLU + +from mlx_streaming import config +from mlx_streaming.core.moe.custom_kernel import ( + _custom_qproj_enabled, _custom_fused_moe_enabled, _custom_qproj_targets, + _custom_qlinear_indexed, _custom_fused_moe_indexed) + + +def _unique_and_local(inds: mx.array) -> Tuple[mx.array, mx.array]: + """返回 (uniq, local)：uniq 是排序后的唯一全局专家 id，local 是与 inds 同形状的本地下标。""" + flat = [int(i) for i in inds.reshape(-1).tolist()] + uniq_sorted = sorted(set(flat)) + remap = {g: i for i, g in enumerate(uniq_sorted)} + local = mx.array([remap[i] for i in flat], dtype=inds.dtype).reshape(inds.shape) + uniq = mx.array(uniq_sorted, dtype=inds.dtype) + return uniq, local + + +def _slice_switch_linear(lin, uniq: mx.array): + """把一个 SwitchLinear/QuantizedSwitchLinear 沿专家维切到 uniq，返回新的小 linear。""" + n = int(uniq.shape[0]) + has_bias = "bias" in lin + if isinstance(lin, QuantizedSwitchLinear): + new = QuantizedSwitchLinear( + lin.input_dims, lin.output_dims, n, bias=has_bias, + group_size=lin.group_size, bits=lin.bits, mode=lin.mode, + ) + else: + new = SwitchLinear(lin.input_dims, lin.output_dims, n, bias=has_bias) + + sliced = {} + for name, p in lin.parameters().items(): + if isinstance(p, mx.array) and p.ndim >= 1 and p.shape[0] == lin.num_experts: + sliced[name] = p[uniq] # 第 0 维是专家维，按 uniq 取子集 + else: + sliced[name] = p + new.update(sliced) + return new + + +def streaming_switch_glu_forward(glu: SwitchGLU, x: mx.array, inds: mx.array) -> mx.array: + """只在 inds 涉及到的唯一专家上做 SwitchGLU 等价计算。""" + uniq, local = _unique_and_local(inds) + sub = SwitchGLU(glu.gate_proj.input_dims, glu.gate_proj.output_dims, int(uniq.shape[0])) + sub.gate_proj = _slice_switch_linear(glu.gate_proj, uniq) + sub.up_proj = _slice_switch_linear(glu.up_proj, uniq) + sub.down_proj = _slice_switch_linear(glu.down_proj, uniq) + sub.activation = glu.activation + return sub(x, local) + + +def _build_qsl(prefix: str, fetched: dict, in_dims: int, out_dims: int, n: int, + group_size: int, bits: int): + """用 fetched 里以 prefix.* 为键的堆叠参数构造一个 QuantizedSwitchLinear。""" + lin = QuantizedSwitchLinear(in_dims, out_dims, n, bias=False, + group_size=group_size, bits=bits, mode="affine") + sub = {name.split(".", 1)[1]: v for name, v in fetched.items() + if name.startswith(prefix + ".")} + lin.update(sub) + return lin + + +def streaming_switch_glu_forward_from_store(store, layer, x, inds, hidden, moe_inter, + group_size, bits): + """文件后端版：从 store 只取选中专家、构造小 SwitchGLU 计算（等价于全量 SwitchGLU）。 + + 无状态版本（每次新建 sub），保留给单测使用；端到端走 PersistentSubGLU。 + """ + uniq, local = _unique_and_local(inds) + fetched = store.fetch(layer, [int(i) for i in uniq.tolist()]) + n = int(uniq.shape[0]) + sub = SwitchGLU(hidden, moe_inter, n) + sub.gate_proj = _build_qsl("gate_proj", fetched, hidden, moe_inter, n, group_size, bits) + sub.up_proj = _build_qsl("up_proj", fetched, hidden, moe_inter, n, group_size, bits) + sub.down_proj = _build_qsl("down_proj", fetched, moe_inter, hidden, n, group_size, bits) + return sub(x, local) + + +def _update_qsl(lin, prefix: str, fetched: dict): + """原地更新一个 QuantizedSwitchLinear 的参数（只换数组引用，不重建对象）。""" + sub = {name.split(".", 1)[1]: v for name, v in fetched.items() + if name.startswith(prefix + ".")} + lin.update(sub) + + +class PersistentSubGLU: + """按专家数 n 缓存一套 SwitchGLU + 3×QuantizedSwitchLinear，跨调用复用。 + + 解码 batch=1 时每层唯一专家数恒为 top_k，n 不变 → 只在首次/n 变化时构造一次， + 之后每个 token 仅原地 update 三组权重，省掉每调用重建对象与初始化随机量化权重的开销。 + """ + + def __init__(self, hidden: int, moe_inter: int, group_size: int, bits: int, + proj_bits: dict | None = None, layer_idx: int | None = None, + quant_mode: str = "affine", swiglu_limit: float = 0.0): + self.hidden = hidden + self.moe_inter = moe_inter + self.group_size = group_size + self.bits = bits + self.layer_idx = layer_idx + # 混合精度：每个 proj 用各自 bit；None 时三 proj 统一用 bits。 + self.proj_bits = proj_bits or { + "gate_proj": bits, "up_proj": bits, "down_proj": bits} + # quant_mode：量化模式（affine / mxfp4）；swiglu_limit>0 时走 DeepSeek 手动 clip 路径。 + self.quant_mode = quant_mode + self.swiglu_limit = swiglu_limit + self._glu = None + self._n = None + + def _ensure(self, n: int): + if self._glu is not None and self._n == n: + return + pb = self.proj_bits + glu = SwitchGLU(self.hidden, self.moe_inter, n) + glu.gate_proj = QuantizedSwitchLinear( + self.hidden, self.moe_inter, n, bias=False, + group_size=self.group_size, bits=pb["gate_proj"], mode=self.quant_mode) + glu.up_proj = QuantizedSwitchLinear( + self.hidden, self.moe_inter, n, bias=False, + group_size=self.group_size, bits=pb["up_proj"], mode=self.quant_mode) + glu.down_proj = QuantizedSwitchLinear( + self.moe_inter, self.hidden, n, bias=False, + group_size=self.group_size, bits=pb["down_proj"], mode=self.quant_mode) + self._glu = glu + self._n = n + + def forward(self, fetched: dict, n: int, x: mx.array, local: mx.array) -> mx.array: + self._ensure(n) + _update_qsl(self._glu.gate_proj, "gate_proj", fetched) + _update_qsl(self._glu.up_proj, "up_proj", fetched) + _update_qsl(self._glu.down_proj, "down_proj", fetched) + if self.swiglu_limit > 0: + # DeepSeek 路径：手动 gate/up/clip/silu/down，不能用 SwiGLU 融合激活 + # （镜像 deepseek_v4.Experts：up 双侧 clip、gate 上侧 minimum）。 + x_exp = mx.expand_dims(x, (-2, -3)) # (..., 1, 1, D) + gate = self._glu.gate_proj(x_exp, local) + up = self._glu.up_proj(x_exp, local) + up = mx.clip(up, -self.swiglu_limit, self.swiglu_limit) + gate = mx.minimum(gate, self.swiglu_limit) + h = nn.silu(gate) * up + y = self._glu.down_proj(h, local) + return y.squeeze(-2) + max_fused_seq = config.custom_fused_moe_max_seq() + if (x.shape[1] <= max_fused_seq + and _custom_fused_moe_enabled(self.layer_idx or -1, self.proj_bits)): + return self._custom_fused_forward(x, local) + max_seq = config.custom_qproj_max_seq() + if (x.shape[1] <= max_seq + and _custom_qproj_enabled(self.layer_idx or -1, self.proj_bits["gate_proj"])): + return self._custom_gate_up_forward(x, local) + return self._glu(x, local) + + def _custom_fused_forward(self, x: mx.array, local: mx.array) -> mx.array: + """完整替换 gate/up/SwiGLU/down，保持 SwitchGLU 的 [B,S,K,H] 输出契约。""" + shape = local.shape + k = int(shape[-1]) + x_flat = mx.broadcast_to(mx.expand_dims(x, -2), x.shape[:-1] + (k, x.shape[-1])) + x_flat = x_flat.reshape(-1, self.hidden).astype(mx.float32) + idx = local.reshape(-1).astype(mx.uint32) + y = _custom_fused_moe_indexed( + x_flat, idx, + self._glu.gate_proj["weight"], self._glu.gate_proj["scales"], self._glu.gate_proj["biases"], + self._glu.up_proj["weight"], self._glu.up_proj["scales"], self._glu.up_proj["biases"], + self._glu.down_proj["weight"], self._glu.down_proj["scales"], self._glu.down_proj["biases"], + self.hidden, self.moe_inter, self.group_size, self.proj_bits["gate_proj"], + ) + return y.reshape(shape + (self.hidden,)) + + def _custom_gate_up_forward(self, x: mx.array, local: mx.array) -> mx.array: + """只替换 gate/up projection；down 仍走 MLX QuantizedSwitchLinear。""" + shape = local.shape + k = int(shape[-1]) + x_flat = mx.broadcast_to(mx.expand_dims(x, -2), x.shape[:-1] + (k, x.shape[-1])) + x_flat = x_flat.reshape(-1, self.hidden).astype(mx.float32) + idx = local.reshape(-1).astype(mx.uint32) + tile = config.custom_qproj_tile() + up = _custom_qlinear_indexed( + x_flat, idx, + self._glu.up_proj["weight"], self._glu.up_proj["scales"], self._glu.up_proj["biases"], + self.moe_inter, self.hidden, self.group_size, self.proj_bits["up_proj"], tile) + gate = _custom_qlinear_indexed( + x_flat, idx, + self._glu.gate_proj["weight"], self._glu.gate_proj["scales"], self._glu.gate_proj["biases"], + self.moe_inter, self.hidden, self.group_size, self.proj_bits["gate_proj"], tile) + up = up.reshape(shape + (1, self.moe_inter)) + gate = gate.reshape(shape + (1, self.moe_inter)) + a = self._glu.activation(up, gate) + if "down" in _custom_qproj_targets(): + a_flat = a.squeeze(-2).reshape(-1, self.moe_inter).astype(mx.float32) + y = _custom_qlinear_indexed( + a_flat, idx, + self._glu.down_proj["weight"], self._glu.down_proj["scales"], self._glu.down_proj["biases"], + self.hidden, self.moe_inter, self.group_size, self.proj_bits["down_proj"], tile) + return y.reshape(shape + (self.hidden,)) + y = self._glu.down_proj(a, local) + return y.squeeze(-2) diff --git a/mlx_streaming/core/moe/custom_kernel.py b/mlx_streaming/core/moe/custom_kernel.py new file mode 100644 index 0000000..2bc6eeb --- /dev/null +++ b/mlx_streaming/core/moe/custom_kernel.py @@ -0,0 +1,266 @@ +"""自定义 Metal 算子：indexed 量化线性 + fused MoE expert（实验性加速路径）。 + +这些 kernel 走 `mx.fast.metal_kernel`，按 `indices[pair]` 选择专家权重， +直接在 packed 量化权重上算 SwiGLU/down，避免反量化中间张量。仅在对应 +CUSTOM_* 环境开关命中且 bit 宽匹配时启用，默认关闭（不改变模型数值）。 +""" +from typing import Tuple # noqa: F401 (保留以兼容历史 import) + +import mlx.core as mx + +from mlx_streaming import config +from mlx_streaming.config import parse_layers_env as _parse_layers_env + +# kernel 编译缓存：按 (tile) / (hidden,moe_inter,...) 复用已编译 kernel。 +_CUSTOM_QKERNEL_CACHE = {} +_CUSTOM_FUSED_MOE_CACHE = {} + + +def _custom_qproj_enabled(layer_idx: int, bits: int) -> bool: + if not config.custom_qproj(): + return False + if bits != config.custom_qproj_bits(): + return False + layers = _parse_layers_env("CUSTOM_QPROJ_LAYERS") + return layers is None or int(layer_idx) in layers + + +def _custom_fused_moe_enabled(layer_idx: int, proj_bits: dict) -> bool: + if not config.custom_fused_moe(): + return False + bits = config.custom_fused_moe_bits() + if any(int(proj_bits[name]) != bits for name in ("gate_proj", "up_proj", "down_proj")): + return False + layers = _parse_layers_env("CUSTOM_FUSED_MOE_LAYERS") + return layers is None or int(layer_idx) in layers + + +def _custom_qproj_targets() -> set[str]: + spec = config.custom_qproj_targets() + return {x.strip() for x in spec.split(",") if x.strip()} + + +def _custom_qlinear_indexed(x: mx.array, indices: mx.array, weight: mx.array, + scales: mx.array, biases: mx.array, + out_dim: int, in_dim: int, group_size: int, + bits: int, tile: int = 4) -> mx.array: + """indexed custom qlinear: x[p,in] 使用 indices[p] 选择专家权重，输出 [p,out]。""" + key = tile + kernel = _CUSTOM_QKERNEL_CACHE.get(key) + if kernel is None: + source = r""" + constexpr int lanes_per_row = 256 / rows_per_group; + constexpr int block_size = rows_per_group * lanes_per_row; + uint tid = thread_position_in_threadgroup.x; + uint group_id = thread_position_in_grid.x / block_size; + uint local_row = tid / lanes_per_row; + uint lane = tid % lanes_per_row; + uint global_row = group_id * rows_per_group + local_row; + if (global_row >= pairs * out_dim) return; + uint pair = global_row / out_dim; + uint out_row = global_row % out_dim; + uint expert = indices[pair]; + constexpr uint mask = (1u << bits) - 1u; + constexpr int words_per_row = (in_dim * bits) / 32; + constexpr int groups_per_row = in_dim / group_size; + threadgroup float partial[block_size]; + float acc = 0.0f; + for (int col = int(lane); col < in_dim; col += lanes_per_row) { + int bit_offset = col * bits; + int word_idx = bit_offset / 32; + int shift = bit_offset % 32; + uint base = (expert * out_dim + out_row) * words_per_row; + uint word = weight[base + word_idx]; + uint q = (word >> shift); + if (shift + bits > 32) { + uint next_word = weight[base + word_idx + 1]; + q |= (next_word << (32 - shift)); + } + q = q & mask; + int g = col / group_size; + uint sb = (expert * out_dim + out_row) * groups_per_row + g; + float wv = float(q) * scales[sb] + biases[sb]; + acc += wv * x[pair * in_dim + col]; + } + partial[tid] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (lane < stride) { + partial[tid] += partial[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lane == 0) { + y[global_row] = partial[0 + (local_row * lanes_per_row)]; + } + """ + kernel = mx.fast.metal_kernel( + name=f"custom_qlinear_indexed_tile{tile}", + input_names=["x", "indices", "weight", "scales", "biases"], + output_names=["y"], + source=source, + ) + _CUSTOM_QKERNEL_CACHE[key] = kernel + pairs = int(indices.size) + grid_groups = (pairs * out_dim + tile - 1) // tile + (y,) = kernel( + inputs=[x, indices.astype(mx.uint32), weight, scales, biases], + output_shapes=[(pairs, out_dim)], + output_dtypes=[mx.float32], + grid=(grid_groups * 256, 1, 1), + threadgroup=(256, 1, 1), + template=[ + ("pairs", pairs), + ("out_dim", out_dim), + ("in_dim", in_dim), + ("group_size", group_size), + ("bits", bits), + ("rows_per_group", tile), + ], + ) + return y + + +def _custom_fused_moe_indexed(x: mx.array, indices: mx.array, + gate_w: mx.array, gate_s: mx.array, gate_b: mx.array, + up_w: mx.array, up_s: mx.array, up_b: mx.array, + down_w: mx.array, down_s: mx.array, down_b: mx.array, + hidden: int, moe_inter: int, group_size: int, + bits: int) -> mx.array: + """fused MoE expert: x[p] 用 indices[p] 选择专家，输出 down 后的 hidden。""" + lanes_per_row = config.custom_fused_moe_lanes() + block_size = config.custom_fused_moe_block() + key = (hidden, moe_inter, group_size, bits, lanes_per_row, block_size) + kernel = _CUSTOM_FUSED_MOE_CACHE.get(key) + if kernel is None: + source = r""" + uint tid = thread_position_in_threadgroup.x; + uint pair = thread_position_in_grid.x / block_size; + if (pair >= pairs) return; + uint expert = indices[pair]; + constexpr uint rows_per_step = block_size / lanes_per_row; + uint local_row = tid / lanes_per_row; + uint row_lane = tid % lanes_per_row; + constexpr uint mask = (1u << bits) - 1u; + constexpr int gu_words_per_row = (hidden * bits) / 32; + constexpr int gu_groups_per_row = hidden / group_size; + constexpr int down_words_per_row = (moe_inter * bits) / 32; + constexpr int down_groups_per_row = moe_inter / group_size; + threadgroup float act[1024]; + threadgroup float gate_part[block_size]; + threadgroup float up_part[block_size]; + + // 多个 lane 协作计算同一 row，避免每个 dot 完全串行。 + for (uint row_base = 0; row_base < moe_inter; row_base += rows_per_step) { + uint row = row_base + local_row; + float gate_acc = 0.0f; + float up_acc = 0.0f; + if (row < moe_inter) { + for (uint col = row_lane; col < hidden; col += lanes_per_row) { + int bit_offset = int(col * bits); + int word_idx = bit_offset / 32; + int shift = bit_offset % 32; + uint gu_base = (expert * moe_inter + row) * gu_words_per_row; + uint qg = gate_w[gu_base + word_idx] >> shift; + uint qu = up_w[gu_base + word_idx] >> shift; + if (shift + bits > 32) { + qg |= gate_w[gu_base + word_idx + 1] << (32 - shift); + qu |= up_w[gu_base + word_idx + 1] << (32 - shift); + } + qg &= mask; + qu &= mask; + uint g = col / group_size; + uint sb = (expert * moe_inter + row) * gu_groups_per_row + g; + float xv = x[pair * hidden + col]; + gate_acc += (float(qg) * gate_s[sb] + gate_b[sb]) * xv; + up_acc += (float(qu) * up_s[sb] + up_b[sb]) * xv; + } + } + gate_part[tid] = gate_acc; + up_part[tid] = up_acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (row_lane < stride) { + gate_part[tid] += gate_part[tid + stride]; + up_part[tid] += up_part[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row < moe_inter) { + float gate_v = gate_part[tid]; + float up_v = up_part[tid]; + float sig = 1.0f / (1.0f + exp(-gate_v)); + act[row] = gate_v * sig * up_v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + // down projection 输出最终 hidden。 + for (uint row_base = 0; row_base < hidden; row_base += rows_per_step) { + uint row = row_base + local_row; + float acc = 0.0f; + if (row < hidden) { + for (uint col = row_lane; col < moe_inter; col += lanes_per_row) { + int bit_offset = int(col * bits); + int word_idx = bit_offset / 32; + int shift = bit_offset % 32; + uint base = (expert * hidden + row) * down_words_per_row; + uint q = down_w[base + word_idx] >> shift; + if (shift + bits > 32) { + q |= down_w[base + word_idx + 1] << (32 - shift); + } + q &= mask; + uint g = col / group_size; + uint sb = (expert * hidden + row) * down_groups_per_row + g; + acc += (float(q) * down_s[sb] + down_b[sb]) * act[col]; + } + } + gate_part[tid] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (row_lane < stride) { + gate_part[tid] += gate_part[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row < hidden) { + y[pair * hidden + row] = gate_part[tid]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + """ + kernel = mx.fast.metal_kernel( + name=f"custom_fused_moe_h{hidden}_i{moe_inter}_b{bits}_l{lanes_per_row}", + input_names=[ + "x", "indices", + "gate_w", "gate_s", "gate_b", + "up_w", "up_s", "up_b", + "down_w", "down_s", "down_b", + ], + output_names=["y"], + source=source, + ) + _CUSTOM_FUSED_MOE_CACHE[key] = kernel + pairs = int(indices.size) + (y,) = kernel( + inputs=[ + x, indices.astype(mx.uint32), + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b, + ], + output_shapes=[(pairs, hidden)], + output_dtypes=[mx.float32], + grid=(pairs * block_size, 1, 1), + threadgroup=(block_size, 1, 1), + template=[ + ("pairs", pairs), + ("hidden", hidden), + ("moe_inter", moe_inter), + ("group_size", group_size), + ("bits", bits), + ("lanes_per_row", lanes_per_row), + ("block_size", block_size), + ], + ) + return y diff --git a/mlx_streaming/core/moe/gate.py b/mlx_streaming/core/moe/gate.py new file mode 100644 index 0000000..14fafdf --- /dev/null +++ b/mlx_streaming/core/moe/gate.py @@ -0,0 +1,32 @@ +"""MoE 门控与专家选择：top-k 激活数开关 + 跨层专家预测（gate 前向 + argpartition）。""" +import mlx.core as mx + +from mlx_streaming import config + + +def _effective_top_k(default_k: int) -> int: + """实验开关：降低每 token 激活专家数以测速度/质量取舍。默认不改变模型。""" + override = config.moe_topk_override() + if not override: + return default_k + return max(1, min(default_k, int(override))) + + +def _predict_layer_experts(norm, gate, top_k: int, x: mx.array, mult: int) -> "tuple[dict[int, float], int]": + """用 gate(norm(x)) 预测专家集，返回 ({expert_id: 最大 softmax 分数}, num_experts)。 + + norm/gate 必须取自**目标层**（被预测的那层），以匹配 probe 验证的 + gate_L(post_attention_layernorm_L(h)) 配置（recall_miss≈0.95）。 + """ + gates = mx.softmax(gate(norm(x)), axis=-1, precise=True) + num_experts = gates.shape[-1] + k = min(num_experts, top_k * mult) + inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:] + vals = mx.take_along_axis(gates, inds, axis=-1) + mx.eval(inds, vals) + best: "dict[int, float]" = {} + for e, s in zip(inds.reshape(-1).tolist(), vals.reshape(-1).tolist()): + e, s = int(e), float(s) + if s > best.get(e, -1.0): + best[e] = s + return best, num_experts diff --git a/mlx_streaming/core/moe/native_moe.py b/mlx_streaming/core/moe/native_moe.py new file mode 100644 index 0000000..21a7289 --- /dev/null +++ b/mlx_streaming/core/moe/native_moe.py @@ -0,0 +1,504 @@ +"""CPython native MoE backend 的 Python 包装层。""" +import importlib +import json +import os +import time +from collections import OrderedDict +from functools import lru_cache + +import mlx.core as mx +import numpy as np + +from mlx_streaming import config + + +_EXPERT_STAGE_CACHE: "OrderedDict[tuple, tuple[mx.array, mx.array, mx.array]]" = OrderedDict() +_BUNDLE_STAGE_CACHE: "OrderedDict[tuple, tuple[mx.array, ...]]" = OrderedDict() +_SLOT_POOLS: dict[tuple, "NativeComputeSlotPool"] = {} +_STAGE_STATS = { + "expert_hits": 0, + "expert_misses": 0, + "bundle_hits": 0, + "bundle_misses": 0, + "evictions": 0, + "can_checks": 0, + "can_true": 0, + "can_false": 0, + "calls": 0, + "route_sync_s": 0.0, + "stage_s": 0.0, + "enqueue_s": 0.0, +} + + +class NativeComputeSlotPool: + """按层维护 compute-buffer 常驻 slot，hit 时复用 [cap,...] MLX arrays。""" + + def __init__(self, compute_dir: str, layer: int, hidden: int, inter: int, + group: int, bits: int, num_experts: int, cap: int): + self.compute_dir = compute_dir + self.layer = int(layer) + self.hidden = int(hidden) + self.inter = int(inter) + self.group = int(group) + self.bits = int(bits) + self.num_experts = int(num_experts) + self.cap = int(cap) + self.slot_of: "OrderedDict[int, int]" = OrderedDict() + self.expert_of: list[int | None] = [None] * self.cap + self.slot_arrays = [self._empty_slot() for _ in range(self.cap)] + self.pool_arrays: tuple[mx.array, ...] | None = None + self.dirty = True + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.rebuilds = 0 + + def _empty_slot(self) -> tuple[mx.array, ...]: + gu_words = self.hidden * self.bits // 32 + gu_groups = self.hidden // self.group + down_words = self.inter * self.bits // 32 + down_groups = self.inter // self.group + return ( + mx.zeros((self.inter, gu_words), dtype=mx.uint32), + mx.zeros((self.inter, gu_groups), dtype=mx.uint16), + mx.zeros((self.inter, gu_groups), dtype=mx.uint16), + mx.zeros((self.inter, gu_words), dtype=mx.uint32), + mx.zeros((self.inter, gu_groups), dtype=mx.uint16), + mx.zeros((self.inter, gu_groups), dtype=mx.uint16), + mx.zeros((self.hidden, down_words), dtype=mx.uint32), + mx.zeros((self.hidden, down_groups), dtype=mx.uint16), + mx.zeros((self.hidden, down_groups), dtype=mx.uint16), + ) + + def _load_slot(self, expert: int) -> tuple[mx.array, ...]: + gate = _stage_one_projection_uncached( + self.compute_dir, self.layer, "gate_proj", expert, + self.inter, self.hidden, self.group, self.bits, self.num_experts) + up = _stage_one_projection_uncached( + self.compute_dir, self.layer, "up_proj", expert, + self.inter, self.hidden, self.group, self.bits, self.num_experts) + down = _stage_one_projection_uncached( + self.compute_dir, self.layer, "down_proj", expert, + self.hidden, self.inter, self.group, self.bits, self.num_experts) + return (*gate, *up, *down) + + def _assign_slot(self, expert: int) -> int: + cached = self.slot_of.get(expert) + if cached is not None: + self.slot_of.move_to_end(expert) + self.hits += 1 + return cached + self.misses += 1 + if len(self.slot_of) < self.cap: + slot = next(i for i, e in enumerate(self.expert_of) if e is None) + else: + old_expert, slot = self.slot_of.popitem(last=False) + self.expert_of[slot] = None + self.evictions += 1 + self.slot_arrays[slot] = self._load_slot(expert) + self.expert_of[slot] = expert + self.slot_of[expert] = slot + self.dirty = True + return slot + + def _rebuild_pool_arrays(self) -> None: + self.pool_arrays = tuple( + mx.stack([slot[i] for slot in self.slot_arrays], axis=0) + for i in range(9) + ) + self.rebuilds += 1 + self.dirty = False + + def acquire(self, expert_ids: list[int]) -> tuple[list[int], tuple[mx.array, ...]]: + local = [self._assign_slot(int(e)) for e in expert_ids] + if self.pool_arrays is None or self.dirty: + self._rebuild_pool_arrays() + assert self.pool_arrays is not None + return local, self.pool_arrays + + def stats(self) -> dict: + return { + "cap": self.cap, + "entries": len(self.slot_of), + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "rebuilds": self.rebuilds, + } + + +def _compute_dir() -> str: + return config.compute_buffer_dir() or os.path.join( + config.expert_dir(default=""), "compute_buffers") + + +@lru_cache(maxsize=1) +def _load_ext(): + try: + return importlib.import_module("mlx_streaming.native_moe_ext") + except Exception: + return None + + +@lru_cache(maxsize=256) +def _has_layer_buffers(compute_dir: str, layer: int) -> bool: + for proj in ("gate_proj", "up_proj", "down_proj"): + for name in ("weight", "scales", "biases"): + path = os.path.join(compute_dir, f"layer{layer:02d}.{proj}.{name}.bin") + if not os.path.exists(path): + return False + return True + + +@lru_cache(maxsize=1024) +def _buffer_index_matches(compute_dir: str, layer: int, proj: str, + out_dim: int, in_dim: int, group: int, bits: int, + num_experts: int) -> bool: + path = os.path.join(compute_dir, f"layer{layer:02d}.{proj}.index.json") + if not os.path.exists(path): + return False + try: + with open(path) as f: + index = json.load(f) + except Exception: + return False + tensors = index.get("tensors", {}) + weight_shape = tensors.get("weight", {}).get("shape_per_expert") + scale_shape = tensors.get("scales", {}).get("shape_per_expert") + bias_shape = tensors.get("biases", {}).get("shape_per_expert") + expected_weight = [out_dim, in_dim * bits // 32] + expected_grouped = [out_dim, in_dim // group] + return ( + int(index.get("num_experts", -1)) == int(num_experts) + and weight_shape == expected_weight + and scale_shape == expected_grouped + and bias_shape == expected_grouped + ) + + +@lru_cache(maxsize=512) +def _has_matching_compute_buffers(compute_dir: str, layer: int, + hidden: int, inter: int, group: int, + bits: int, num_experts: int) -> bool: + if not _has_layer_buffers(compute_dir, layer): + return False + # index.json 是防呆线：bit/group 或目录错配时宁可回退 MLX，不冒险读错 mmap shape。 + return ( + _buffer_index_matches( + compute_dir, layer, "gate_proj", inter, hidden, group, bits, num_experts) + and _buffer_index_matches( + compute_dir, layer, "up_proj", inter, hidden, group, bits, num_experts) + and _buffer_index_matches( + compute_dir, layer, "down_proj", hidden, inter, group, bits, num_experts) + ) + + +@lru_cache(maxsize=256) +def _projection_memmaps(compute_dir: str, layer: int, proj: str, + out_dim: int, in_dim: int, group: int, bits: int, + num_experts: int): + words = in_dim * bits // 32 + groups = in_dim // group + base = os.path.join(compute_dir, f"layer{layer:02d}.{proj}") + weight = np.memmap( + base + ".weight.bin", + dtype=np.uint32, + mode="r", + shape=(num_experts, out_dim, words), + ) + scales = np.memmap( + base + ".scales.bin", + dtype=np.uint16, + mode="r", + shape=(num_experts, out_dim, groups), + ) + biases = np.memmap( + base + ".biases.bin", + dtype=np.uint16, + mode="r", + shape=(num_experts, out_dim, groups), + ) + return weight, scales, biases + + +def _stage_projection(compute_dir: str, layer: int, proj: str, + expert_ids: list[int], out_dim: int, in_dim: int, + group: int, bits: int, num_experts: int): + if config.native_moe_stage_cache(): + staged = [ + _stage_one_projection( + compute_dir, layer, proj, int(e), out_dim, in_dim, group, bits, num_experts) + for e in expert_ids + ] + return tuple(mx.stack([entry[i] for entry in staged], axis=0) for i in range(3)) + return _stage_projection_uncached( + compute_dir, layer, proj, expert_ids, out_dim, in_dim, group, bits, num_experts) + + +def _stage_projection_uncached(compute_dir: str, layer: int, proj: str, + expert_ids: list[int], out_dim: int, in_dim: int, + group: int, bits: int, num_experts: int): + weight, scales, biases = _projection_memmaps( + compute_dir, layer, proj, out_dim, in_dim, group, bits, num_experts) + ids = np.asarray(expert_ids, dtype=np.int64) + # advanced indexing 已经产出 numpy 临时数组；mx.array 再接管成 MLX-managed staging arrays。 + return ( + mx.array(np.asarray(weight[ids])), + mx.array(np.asarray(scales[ids])), + mx.array(np.asarray(biases[ids])), + ) + + +def _stage_one_projection_uncached(compute_dir: str, layer: int, proj: str, expert: int, + out_dim: int, in_dim: int, group: int, bits: int, + num_experts: int): + weight, scales, biases = _projection_memmaps( + compute_dir, layer, proj, out_dim, in_dim, group, bits, num_experts) + return ( + mx.array(np.asarray(weight[int(expert)])), + mx.array(np.asarray(scales[int(expert)])), + mx.array(np.asarray(biases[int(expert)])), + ) + + +def _stage_one_projection(compute_dir: str, layer: int, proj: str, expert: int, + out_dim: int, in_dim: int, group: int, bits: int, + num_experts: int): + key = (compute_dir, int(layer), proj, int(expert), out_dim, in_dim, group, bits, num_experts) + cached = _EXPERT_STAGE_CACHE.get(key) + if cached is not None: + _EXPERT_STAGE_CACHE.move_to_end(key) + _STAGE_STATS["expert_hits"] += 1 + return cached + _STAGE_STATS["expert_misses"] += 1 + staged = _stage_one_projection_uncached( + compute_dir, layer, proj, expert, out_dim, in_dim, group, bits, num_experts) + _EXPERT_STAGE_CACHE[key] = staged + _EXPERT_STAGE_CACHE.move_to_end(key) + limit = max(0, config.native_moe_stage_cache_experts()) + while limit and len(_EXPERT_STAGE_CACHE) > limit: + _EXPERT_STAGE_CACHE.popitem(last=False) + _STAGE_STATS["evictions"] += 1 + return staged + + +def _slot_pool(compute_dir: str, layer: int, hidden: int, inter: int, group: int, + bits: int, num_experts: int) -> NativeComputeSlotPool: + cap = max(1, config.native_moe_slot_cap()) + key = (compute_dir, int(layer), hidden, inter, group, bits, num_experts, cap) + pool = _SLOT_POOLS.get(key) + if pool is None: + pool = NativeComputeSlotPool( + compute_dir, int(layer), hidden, inter, group, bits, num_experts, cap) + _SLOT_POOLS[key] = pool + return pool + + +def _stage_bundle(compute_dir: str, layer: int, expert_ids: list[int], + hidden: int, inter: int, group: int, bits: int, + num_experts: int): + key = ( + compute_dir, int(layer), tuple(int(e) for e in expert_ids), + hidden, inter, group, bits, num_experts, + ) + if config.native_moe_stage_bundle_cache(): + cached = _BUNDLE_STAGE_CACHE.get(key) + if cached is not None: + _BUNDLE_STAGE_CACHE.move_to_end(key) + _STAGE_STATS["bundle_hits"] += 1 + return cached + _STAGE_STATS["bundle_misses"] += 1 + gate = _stage_projection( + compute_dir, int(layer), "gate_proj", expert_ids, inter, hidden, group, bits, num_experts) + up = _stage_projection( + compute_dir, int(layer), "up_proj", expert_ids, inter, hidden, group, bits, num_experts) + down = _stage_projection( + compute_dir, int(layer), "down_proj", expert_ids, hidden, inter, group, bits, num_experts) + staged = (*gate, *up, *down) + if config.native_moe_stage_bundle_cache(): + _BUNDLE_STAGE_CACHE[key] = staged + _BUNDLE_STAGE_CACHE.move_to_end(key) + limit = max(0, config.native_moe_stage_cache_bundles()) + while limit and len(_BUNDLE_STAGE_CACHE) > limit: + _BUNDLE_STAGE_CACHE.popitem(last=False) + _STAGE_STATS["evictions"] += 1 + return staged + + +def prefetch_native_moe_stage(layer: int, expert_ids: list[int], + hidden: int, inter: int, group: int, bits: int, + num_experts: int) -> bool: + """提前把预测专家拷入 MLX-managed staging cache，供后续 fused native op 复用。""" + if not config.native_moe(): + return False + if not config.native_moe_stage_prefetch(): + return False + if config.native_moe_synthetic(): + return False + compute_dir = _compute_dir() + if not compute_dir or not _has_matching_compute_buffers( + compute_dir, layer, hidden, inter, group, bits, num_experts): + return False + try: + uniq = list(dict.fromkeys(int(e) for e in expert_ids)) + if not uniq: + return False + for proj, out_dim, in_dim in ( + ("gate_proj", inter, hidden), + ("up_proj", inter, hidden), + ("down_proj", hidden, inter), + ): + for e in uniq: + _stage_one_projection( + compute_dir, int(layer), proj, e, out_dim, in_dim, group, bits, num_experts) + return True + except Exception: + if config.native_moe_raise(): + raise + return False + + +def stage_cache_stats() -> dict: + slot_stats = { + "pools": len(_SLOT_POOLS), + "entries": sum(len(pool.slot_of) for pool in _SLOT_POOLS.values()), + "hits": sum(pool.hits for pool in _SLOT_POOLS.values()), + "misses": sum(pool.misses for pool in _SLOT_POOLS.values()), + "evictions": sum(pool.evictions for pool in _SLOT_POOLS.values()), + "rebuilds": sum(pool.rebuilds for pool in _SLOT_POOLS.values()), + } + return { + **_STAGE_STATS, + "route_sync_s": round(_STAGE_STATS["route_sync_s"], 6), + "stage_s": round(_STAGE_STATS["stage_s"], 6), + "enqueue_s": round(_STAGE_STATS["enqueue_s"], 6), + "expert_entries": len(_EXPERT_STAGE_CACHE), + "bundle_entries": len(_BUNDLE_STAGE_CACHE), + "slot_pool": slot_stats, + } + + +def note_route_sync(seconds: float) -> None: + _STAGE_STATS["route_sync_s"] += float(seconds) + + +def can_native_moe(layer: int, hidden: int, inter: int, group: int, bits: int, + num_experts: int) -> bool: + """只做缓存化的轻量可用性判断；失败时避免上层触发 route host 同步。""" + _STAGE_STATS["can_checks"] += 1 + if not config.native_moe(): + _STAGE_STATS["can_false"] += 1 + return False + if not config.native_moe_mlx_op(): + _STAGE_STATS["can_false"] += 1 + return False + if config.native_moe_synthetic(): + ok = _load_ext() is not None + _STAGE_STATS["can_true" if ok else "can_false"] += 1 + return ok + if (hidden * bits) % 32 != 0 or (inter * bits) % 32 != 0: + _STAGE_STATS["can_false"] += 1 + return False + compute_dir = _compute_dir() + if not compute_dir: + _STAGE_STATS["can_false"] += 1 + return False + ok = ( + _load_ext() is not None + and _has_matching_compute_buffers( + compute_dir, int(layer), hidden, inter, group, bits, num_experts) + ) + _STAGE_STATS["can_true" if ok else "can_false"] += 1 + return ok + + +def try_native_moe(layer: int, expert_ids: list[int], x: mx.array, scores: mx.array, + hidden: int, inter: int, group: int, bits: int, + num_experts: int) -> "mx.array | None": + """尝试 native fused MoE，失败返回 None 让调用方回退 MLX 路径。""" + if not config.native_moe(): + return None + if not config.native_moe_mlx_op(): + return None + if (hidden * bits) % 32 != 0 or (inter * bits) % 32 != 0: + return None + if not expert_ids: + return None + if not can_native_moe(layer, hidden, inter, group, bits, num_experts): + return None + compute_dir = _compute_dir() + ext = _load_ext() + if ext is None: + return None + synthetic = config.native_moe_synthetic() + if (not synthetic and not _has_matching_compute_buffers( + compute_dir, layer, hidden, inter, group, bits, num_experts)): + return None + try: + expert_arr = mx.array([int(e) for e in expert_ids], dtype=mx.uint32) + if synthetic: + t_enqueue = time.perf_counter() + y = ext.fused_moe( + x.astype(mx.float32), + expert_arr, + scores.astype(mx.float32), + compute_dir, + int(layer), + hidden, + inter, + group, + bits, + num_experts, + True, + ) + _STAGE_STATS["enqueue_s"] += time.perf_counter() - t_enqueue + else: + t_stage = time.perf_counter() + if config.native_moe_slot_pool(): + pool = _slot_pool( + compute_dir, int(layer), hidden, inter, group, bits, num_experts) + local_slots, pool_arrays = pool.acquire(expert_ids) + local_arr = mx.array(local_slots, dtype=mx.uint32) + ( + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b, + ) = pool_arrays + else: + local_arr = None + ( + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b, + ) = _stage_bundle( + compute_dir, int(layer), expert_ids, + hidden, inter, group, bits, num_experts) + _STAGE_STATS["stage_s"] += time.perf_counter() - t_stage + t_enqueue = time.perf_counter() + if local_arr is not None: + y = ext.fused_moe_slots( + x.astype(mx.float32), + local_arr, + scores.astype(mx.float32), + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b, + hidden, inter, group, bits, + ) + else: + y = ext.fused_moe_staged( + x.astype(mx.float32), + scores.astype(mx.float32), + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b, + hidden, inter, group, bits, + ) + _STAGE_STATS["enqueue_s"] += time.perf_counter() - t_enqueue + _STAGE_STATS["calls"] += 1 + return y + except Exception: + if config.native_moe_raise(): + raise + return None diff --git a/mlx_streaming/core/prefetch/__init__.py b/mlx_streaming/core/prefetch/__init__.py new file mode 100644 index 0000000..0e0e1ef --- /dev/null +++ b/mlx_streaming/core/prefetch/__init__.py @@ -0,0 +1 @@ +"""专家预取模块：后台物化(bg_prefetch) / native staging(native_staging) / 跨层预测预取(cross_layer) / 模型 patch(patch)。""" diff --git a/mlx_streaming/core/prefetch/bg_prefetch.py b/mlx_streaming/core/prefetch/bg_prefetch.py new file mode 100644 index 0000000..a0e8503 --- /dev/null +++ b/mlx_streaming/core/prefetch/bg_prefetch.py @@ -0,0 +1,133 @@ +"""后台专家预取器：在独立 MLX stream 上物化预测专家（私有 array），交接给主线程。 + +经 gating 测试验证可行（probe_multistream_gate / probe_multistream_handoff）： +- 后台线程 `with mx.stream(s2):` 物化 + mx.eval，可与主线程计算重叠、不崩。 +- 只物化「私有」array，绝不写主线程共享池（跨 stream 写共享张量会报 no Stream 错）。 +- 主线程只消费「已 eval」的交接 array → 跨 stream 读安全。 +""" +import queue +import threading +from collections import OrderedDict + +import mlx.core as mx + +from mlx_streaming import config + + +class BackgroundExpertPrefetcher: + def __init__(self, blob_source, window: int = 3, native: bool = False): + self._src = blob_source + self._stream = mx.new_stream(mx.default_device()) + self._q: "queue.Queue" = queue.Queue() + self._ready: "OrderedDict[tuple, dict]" = OrderedDict() + self._ready_layers: "list[int]" = [] + self._window = window + self._native = native + self._lock = threading.Lock() + self._stop = False + self.submitted = 0 + self.materialized = 0 + self.taken = 0 + # 就绪率诊断：promote 时已物化好的(ready_on_time) vs 提交了但还在飞行中(not_ready)。 + # 用于量化 attention/GDN 窗口是否够长隐藏 I/O。 + self.ready_on_time = 0 + self.not_ready = 0 + self.materialize_s = 0.0 # bg 线程物化(load+eval)累计墙钟,验证"串行加性"假设 + self._native_mat = config.native_materialize() + self._inflight: "dict[int, set]" = {} + self._t = threading.Thread(target=self._loop, daemon=True) + self._t.start() + + def submit(self, layer: int, expert_ids) -> None: + ids = [int(e) for e in expert_ids] + if not ids: + return + with self._lock: + self.submitted += len(ids) + self._inflight.setdefault(int(layer), set()).update(ids) + self._q.put((int(layer), ids)) + + def _loop(self) -> None: + while not self._stop: + try: + layer, ids = self._q.get(timeout=0.1) + except queue.Empty: + continue + try: + import time as _t + _t0 = _t.perf_counter() + with mx.stream(self._stream): + # view_bf16=False：后台不做 .view(bfloat16)（会报 no Stream(gpu,0)）； + # scales/biases 留 uint16，主线程 take 时再 view。 + # NATIVE_MATERIALIZE=1：用 C++ blob_load 把拷贝挪出 GIL（减少对主线程争用）。 + if self._native_mat: + experts = self._src.load_experts_native(layer, ids, view_bf16=False) + else: + experts = self._src.load_experts(layer, ids, view_bf16=False) + mx.eval([v for d in experts.values() for v in d.values()]) + self.materialize_s += _t.perf_counter() - _t0 + with self._lock: + pend = self._inflight.get(layer) + for e, d in experts.items(): + self._ready[(layer, int(e))] = d + self.materialized += 1 + if pend is not None: + pend.discard(int(e)) + if layer not in self._ready_layers: + self._ready_layers.append(layer) + while len(self._ready_layers) > self._window: + old = self._ready_layers.pop(0) + for k in [k for k in self._ready if k[0] == old]: + del self._ready[k] + except Exception: + # 后台失败不影响主线程：主线程会优雅回退到同步 demand 路径。 + pass + + @staticmethod + def _view_bf16(d: dict) -> dict: + """主线程消费时把 scales/biases 从 uint16 位重解释回 bfloat16。""" + out = {} + for k, v in d.items(): + out[k] = v.view(mx.bfloat16) if (k.endswith(".scales") or k.endswith(".biases")) else v + return out + + def take_ready(self, layer: int, e: int) -> "dict | None": + with self._lock: + d = self._ready.pop((int(layer), int(e)), None) + if d is None: + return None + self.taken += 1 + return self._view_bf16(d) + + def take_ready_layer(self, layer: int) -> "dict[int, dict]": + """取走某层所有就绪专家：{expert_id: {proj.tensor: mx.array}}（主线程调用）。""" + raw = {} + with self._lock: + for key in [k for k in self._ready if k[0] == int(layer)]: + raw[key[1]] = self._ready.pop(key) + self.taken += 1 + return {e: self._view_bf16(d) for e, d in raw.items()} + + def ready_count(self, layer: int) -> int: + with self._lock: + return sum(1 for k in self._ready if k[0] == int(layer)) + + def note_promote(self, layer: int, ready_now: int) -> None: + """promote_prefetched 调用：记录本次 promote 时已就绪数 vs 仍在飞行中(窗口没盖住)。""" + with self._lock: + self.ready_on_time += int(ready_now) + pend = self._inflight.get(int(layer)) + if pend: + self.not_ready += len(pend) + pend.clear() + + def stats(self) -> dict: + with self._lock: + return {"submitted": self.submitted, "materialized": self.materialized, + "taken": self.taken, "ready": len(self._ready), + "ready_on_time": self.ready_on_time, "not_ready": self.not_ready, + "materialize_s": round(self.materialize_s, 3)} + + def close(self) -> None: + self._stop = True + self._t.join(timeout=2) diff --git a/mlx_streaming/core/prefetch/cross_layer.py b/mlx_streaming/core/prefetch/cross_layer.py new file mode 100644 index 0000000..35691b3 --- /dev/null +++ b/mlx_streaming/core/prefetch/cross_layer.py @@ -0,0 +1,181 @@ +"""跨层专家预取：在 attention/GDN 之前用本层输入预测目标层 MoE 路由并提前取专家。 + +通过 monkeypatch `Qwen3NextDecoderLayer.__call__`，在真正进入 MoE 计算前用 +gate_L(post_attention_layernorm_L(h)) 预测目标层专家（对齐 recall≈0.95 的探针口径）， +把 top-(top_k*mult) 专家在当前层计算窗口内异步预取，藏住读盘/物化延迟。 + +支持多条预取后端（按环境开关择一）：STREAM_BLOB_BG（后台物化进池）、STREAM_BLOB +（blob 字节预读）、STREAM_BLOB_LOADER（blob-loader 字节预热）、native stage 预取。 +""" +import time + +import mlx.core as mx + +from mlx_streaming import config +from mlx_streaming.core.moe import native_moe +from mlx_streaming.core.moe.gate import _predict_layer_experts +from mlx_streaming.core.moe.block import FileStreamingMoeBlock + +# 是否已对 decoder layer 打过预取补丁（全局只打一次）。 +_CROSS_LAYER_PREFETCH_PATCHED = False +# 每次「整模型前向」内的全局预取预算（按层序回绕重置），避免无界预取撑爆窗口。 +_FILE_PREFETCH_REMAINING = 0 +_FILE_PREFETCH_LAST_LAYER = -1 +_STAGE_PREFETCH_REMAINING = 0 +_STAGE_PREFETCH_LAST_LAYER = -1 + + +def _cross_layer_prefetch_mult() -> int: + return config.cross_layer_mult() + + +def _submit_missing_prefetch(store, layer: int, predicted) -> "list[int]": + """只把"预测∩非常驻"的专家提交给后台预取器，避免重载常驻、撑爆 attention 窗口。 + + 返回实际提交的缺失专家列表（供测试/统计）。 + """ + bg = getattr(store, "_bg", None) + if bg is None: + return [] + resident = store.resident_experts(layer) if hasattr(store, "resident_experts") else set() + missing = [int(e) for e in predicted if int(e) not in resident] + if missing: + bg.submit(int(layer), missing) + return missing + + +def _prefetch_budget(layer_idx: int, kind: str) -> int: + """按一次层序前向重置全局预算；layer 序号回绕说明进入新 forward。""" + global _FILE_PREFETCH_REMAINING, _FILE_PREFETCH_LAST_LAYER + global _STAGE_PREFETCH_REMAINING, _STAGE_PREFETCH_LAST_LAYER + if kind == "file": + total = config.file_prefetch_global_budget() + if layer_idx <= _FILE_PREFETCH_LAST_LAYER: + _FILE_PREFETCH_REMAINING = total + _FILE_PREFETCH_LAST_LAYER = layer_idx + return max(0, _FILE_PREFETCH_REMAINING) + total = config.stage_prefetch_global_budget() + if layer_idx <= _STAGE_PREFETCH_LAST_LAYER: + _STAGE_PREFETCH_REMAINING = total + _STAGE_PREFETCH_LAST_LAYER = layer_idx + return max(0, _STAGE_PREFETCH_REMAINING) + + +def _consume_prefetch_budget(n: int, kind: str) -> None: + global _FILE_PREFETCH_REMAINING, _STAGE_PREFETCH_REMAINING + if kind == "file": + _FILE_PREFETCH_REMAINING = max(0, _FILE_PREFETCH_REMAINING - n) + else: + _STAGE_PREFETCH_REMAINING = max(0, _STAGE_PREFETCH_REMAINING - n) + + +def _prefetch_layers(kind: str) -> "set[int] | None": + env = "FILE_PREFETCH_LAYERS" if kind == "file" else "STAGE_PREFETCH_LAYERS" + return config.parse_layers_env(env) + + +def enable_cross_layer_prefetch(): + """给 Qwen3NextDecoderLayer 加跨层专家预取。 + + 在 attention/GDN 之前，用当前层输入 hidden 的 post_attention_layernorm 近似预测 + 本层 MoE 路由，并把 top-(top_k*mult) 专家预取进 resident pool。 + """ + global _CROSS_LAYER_PREFETCH_PATCHED + if _CROSS_LAYER_PREFETCH_PATCHED: + return + from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer + orig_call = Qwen3NextDecoderLayer.__call__ + + def patched_call(self, x, mask=None, cache=None): + mlp = getattr(self, "mlp", None) + if mlp is not None and getattr(getattr(mlp, "store", None), "_staging", None) is not None: + # 存本层「未归一化」decoder 输入：供 native-fused-prefetch 用目标层 norm 正确预测 + # （对齐 0.95-recall 探针；MoE forward 拿到的 x 已被本层 norm 过，norm 用错会掉点）。 + mlp._unnormed_input = x + if config.cross_layer_prefetch() and config.resident_pool_enabled(): + ahead = config.cross_layer_ahead(default=0) + target_layer = getattr(self, "_layer_idx", None) + if target_layer is not None: + target_layer += ahead + target_mlp = None + target_decoder = None + if ahead == 0 and isinstance(mlp, FileStreamingMoeBlock): + target_mlp = mlp + target_decoder = self + elif target_layer is not None: + layers = getattr(getattr(self, "_prefetch_model_ref", None), "layers", []) + if 0 <= target_layer < len(layers): + target_decoder = layers[target_layer] + target_mlp = getattr(target_decoder, "mlp", None) + if isinstance(target_mlp, FileStreamingMoeBlock): + if config.probe_predict_only(): + # 诊断：只做预测 gate 前向 + eval(同步在预测上)，不 .tolist、不预取。 + # 用于拆分 32% 税：≈S(+9%)→税是 Python 编排(native 可救)；≈B(+32%)→税是 + # gate前向被 barrier 暴露在关键路径(简单 native 派发救不了)。 + g = target_mlp.gate(target_decoder.post_attention_layernorm(x)) + kk = min(g.shape[-1], 4) + ip = mx.argpartition(g, kth=-kk, axis=-1)[..., -kk:] + mx.eval(ip) + return orig_call(self, x, mask=mask, cache=cache) + # 关键：用**目标层**的 post_attention_layernorm（匹配 probe 验证的 + # gate_L(post_norm_L(h_{L-ahead})) 配置，recall_miss≈0.95）。 + best, num_experts = _predict_layer_experts( + target_decoder.post_attention_layernorm, target_mlp.gate, + target_mlp.top_k, x, _cross_layer_prefetch_mult()) + bg = getattr(target_mlp.store, "_bg", None) + if config.stream_blob_bg() and bg is not None: + # 同层(AHEAD=0)预测 + 只预取"预测∩非常驻"，藏进 attention/GDN 窗口。 + budget = config.stream_blob_bg_budget( + default=target_mlp.top_k * _cross_layer_prefetch_mult()) + picked = [e for e, _ in sorted(best.items(), key=lambda kv: kv[1], reverse=True)][:budget] + _submit_missing_prefetch(target_mlp.store, target_mlp.layer_idx, picked) + if config.window_prof(): + target_mlp._submit_t = time.perf_counter() + return orig_call(self, x, mask=mask, cache=cache) + if config.stream_blob() and target_mlp._blob is not None: + # 全流式：后台预读下一层预测专家的字节，与当前层计算重叠。 + budget = config.stream_blob_prefetch_budget(default=target_mlp.top_k * 2) + picked = [e for e, _ in sorted(best.items(), key=lambda kv: kv[1], reverse=True)][:budget] + target_mlp._blob.prefetch_async(target_mlp.layer_idx, picked) + return orig_call(self, x, mask=mask, cache=cache) + blob_loader = getattr(target_mlp.store, "_blob_loader", None) + if config.stream_blob_loader() and blob_loader is not None: + # blob-loader 路径：后台只预读字节（与当前层计算重叠）， + # demand 时主线程从预热字节快速物化进常驻池，不在主线程同步阻塞读盘。 + budget = config.stage_prefetch_per_layer_budget(default=12) + picked = [e for e, _ in sorted(best.items(), key=lambda kv: kv[1], reverse=True)][:budget] + blob_loader.prefetch_async(target_mlp.layer_idx, picked) + return orig_call(self, x, mask=mask, cache=cache) + allowed_layers = _prefetch_layers("stage") + if allowed_layers is not None and target_mlp.layer_idx not in allowed_layers: + return orig_call(self, x, mask=mask, cache=cache) + min_score = config.stage_prefetch_min_score() + per_layer = config.stage_prefetch_per_layer_budget(default=4) + remaining = _prefetch_budget(target_mlp.layer_idx, "stage") + budget = max(0, min(per_layer, remaining)) + expert_ids = [ + e for e, s in sorted(best.items(), key=lambda kv: kv[1], reverse=True) + if s >= min_score + ][:budget] + _consume_prefetch_budget(len(expert_ids), "stage") + proj_bits = getattr(target_mlp._sub, "proj_bits", {}) + bits_set = { + int(proj_bits.get(name, target_mlp.bits)) + for name in ("gate_proj", "up_proj", "down_proj") + } + if len(bits_set) == 1 and native_moe.prefetch_native_moe_stage( + target_mlp.layer_idx, + expert_ids, + target_mlp.hidden, + target_mlp.moe_inter, + target_mlp.group_size, + bits_set.pop(), + num_experts, + ): + pass + else: + target_mlp.store.prefetch(target_mlp.layer_idx, expert_ids) + return orig_call(self, x, mask=mask, cache=cache) + + Qwen3NextDecoderLayer.__call__ = patched_call + _CROSS_LAYER_PREFETCH_PATCHED = True diff --git a/mlx_streaming/core/prefetch/native_staging.py b/mlx_streaming/core/prefetch/native_staging.py new file mode 100644 index 0000000..9e5a480 --- /dev/null +++ b/mlx_streaming/core/prefetch/native_staging.py @@ -0,0 +1,229 @@ +"""native-fused-prefetch 的 miss→hit 落地：per-layer 环形 staging buffer + 池 promote。 + +机制（预读 IO 全在 GPU 完成回调里、主线程零 pread/零 .tolist/零 np→mx.array）： +- submit(layer, inds_lazy)：取该层环里的一块 buffer，调 prefetch_into_staging → + GPU 完成回调（C++）把专家字节 pread 进这块 buffer，并按 gen 原子记录 (expert→row)。 +- promote(layer, store)：读 C++ 记录（纯锁）→ 按 gen 取回 handler 真正写过的那块 buffer → + 惰性切片直接 _place_expert 写进池槽（不 eval，避免每层 host 同步拖慢 verify 流水线）。 + +为何必须 per-layer（而非全局共享小环）：预读在 GPU 完成回调里异步发生，且 batched verify +是「整块前向、最后一次 eval」——同一前向里全部 48 层的 submit 都在飞行，其 handler 可能要到 +接近末尾才触发。若用全局小环，buffer 会在它的 handler 落地/promote 消费**之前**就被后续层的 +submit 复用、被 handler 覆盖成别层数据（gen 匹配只保证 buffer 对象、保证不了内容新鲜）→ 非确定性 +串扰损坏（实测全局环 n_mismatch=60~85）。所以每层各留独立 buffer，环大小只用于跨 token 复用。 + +内存 = 层数 × ring × budget × stride。ring 默认从 4 降到 2（STAGING_RING），在保持逐位精确的前提下 +把 staging 占用减半（48×2×budget×stride）。 +""" +import os +import time + +import mlx.core as mx + +from mlx_streaming import config +from mlx_streaming.core.profiling import note_tprof, TPROF_ON + +# 诊断门控（STG_VERIFY）：开启时 promote 记录每个 (layer, expert) 落池所用 gen，供消费侧字节校验 +# 定位损坏来自哪一代 staging。默认关，对主路径零影响（关时连 dict 写入都不发生）。 +_STG_VERIFY = os.environ.get("STG_VERIFY") == "1" + + +def _typed_seg(seg, dt, shape): + """按段 dtype 重解释：uint32 原样 / uint16→bfloat16 / uint8(mxfp4 scales) 原样。""" + if dt == "uint32": + return seg.view(mx.uint32).reshape(shape) + if dt == "uint16": + return seg.view(mx.uint16).reshape(shape).view(mx.bfloat16) + return seg.reshape(shape) + + +def route_used_subset(cand: "list[int]", route_inds: mx.array, num_experts: int) -> "set[int]": + """返回 cand 中“被 route_inds 真实路由到”的子集（GPU membership，drain ≤len(cand)）。 + + cand：预读候选专家 id（host 已知，通常 ≤budget）。 + route_inds：本层真实路由 id（GPU lazy，任意 shape，内部 reshape(-1)）。 + num_experts：专家总数（建 mask 用）。 + cand 为空 → 返回空集（不触发任何 GPU op）。 + + 仅对预读候选做成员测试：不物化全量路由，只把 ≤len(cand) 个布尔拉回 host。 + """ + if not cand: + return set() + route_mask = mx.zeros((num_experts,), dtype=mx.uint8) + route_mask[route_inds.reshape(-1)] = 1 # GPU scatter（幂等） + hit = route_mask[mx.array(cand, dtype=mx.uint32)] # GPU gather，≤len(cand) + hit_list = hit.tolist() # 唯一一次极小 drain + return {int(cand[i]) for i in range(len(cand)) if hit_list[i]} + + +class NativeStagingManager: + """per-layer 环形 staging：每层 ring 块 buffer 轮转，按 gen 严格匹配回收。 + + submit 写 ring[layer][rr]、记 (gen→buffer)；C++ handler 按 gen 原子记 (expert→row)； + promote 按 gen 取回**正是 handler 写过**的那块 buffer。只要某层的 buffer 在该层下一次（ring 个 + submit 之后）被复用前，其 handler 已落地且 promote 的惰性切片已随本前向 eval 被消费，即不串扰。 + """ + def __init__(self, blob_source, budget: int, ring: int | None = None): + self.src = blob_source # BlobExpertSource：有 dir / stride / _segs + self.budget = int(budget) + self.stride = int(blob_source.stride) + # 安全下限=2：MTP 每步 verify+replay 会对同层各 submit 一次（2 个在飞），ring=1 实测串扰损坏。 + self.ring = max(2, int(ring) if ring is not None else config.staging_ring()) + self._ring: "dict[int, list]" = {} # layer -> [mx.array]*ring + self._rr: "dict[int, int]" = {} # layer -> 下一个要写的环索引 + self._gen = 1 # 全局单调 generation + self._gen_buf: "dict[int, mx.array]" = {} # gen -> 该次 submit 写的 buffer + self.submitted = 0 + self.promoted = 0 + # 诊断用:每层最近一次 promote take 到的"已就绪(C++ pread 完成)专家集"。 + # miss_attrib 据此把 miss_A 拆成 时序(不在就绪集) / 驱逐(在就绪集但 acquire 前已不在池)。 + self.last_ready: "dict[int, set]" = {} + # 诊断用(STG_VERIFY)：记录每个 (layer, expert) 最近一次被 promote 写池时所用的 gen， + # 供消费侧字节校验定位「损坏来自哪一代 staging」。默认不影响主路径(只是写 dict)。 + self.placed_gen: "dict[tuple, int]" = {} + + def _bufs(self, layer: int) -> list: + bl = self._ring.get(layer) + if bl is None: + bl = [] + for _ in range(self.ring): + a = mx.zeros((self.budget, self.stride), dtype=mx.uint8) + mx.eval(a) + bl.append(a) + self._ring[layer] = bl + self._rr[layer] = 0 + return bl + + def submit(self, layer: int, inds_lazy: mx.array, resident: "list[int] | None" = None): + """方案B：inds_lazy 为"预测宽集合"(lazy uint32 [N]，按门控分降序，N 可远大于 budget)。 + + resident：目标层提交时刻的常驻专家快照（host 侧）。C++ 回调按 resident 过滤、再按降序 + 取前 budget 个"缺口"专家 pread 进 buffer（budget=buffer 行数）。返回 dummy（折进图触发回调）。 + """ + import mlx_streaming.native_moe_ext as _N + layer = int(layer) + bufs = self._bufs(layer) + rr = self._rr[layer] + buf = bufs[rr] + self._rr[layer] = (rr + 1) % self.ring + gen = self._gen + self._gen += 1 + self._gen_buf[gen] = buf + if len(self._gen_buf) > self.ring * 64: + for g in sorted(self._gen_buf)[:len(self._gen_buf) - self.ring * 64]: + self._gen_buf.pop(g, None) + path = f"{self.src.dir}/layer{layer:02d}.blob" + self.submitted += 1 + res = [int(e) for e in (resident or [])] + # cap=self.budget：回调最多往 buffer 写 budget 行（覆盖缺口分布，p99≈15 → budget=16 够）。 + return _N.prefetch_into_staging( + buf, inds_lazy, layer, gen, path, self.stride, res, self.budget, + config.staging_pread_parallel()) + + def promote(self, layer: int, store, used: "set[int] | None" = None, + route_inds: "mx.array | None" = None, + num_experts: "int | None" = None) -> int: + """把该层 staging 里已就绪专家写进常驻池（主线程、acquire 前）。 + + used：本层真实路由专家集合（host 已知时传入）。只 promote "预读好 ∩ used" 的专家。 + route_inds/num_experts：GPU 重映射路径专用——host 没有现成 used 时， + 用 route_used_subset 仅对预读候选做 GPU membership 现算 used（drain ≤budget）。 + 三者皆缺 → 整批写入（兜底，行为同改动前）。 + """ + import mlx_streaming.native_moe_ext as _N + layer = int(layer) + _pt0 = time.perf_counter() if TPROF_ON else 0.0 + flat = _N.prefetch_staging_take(layer) + if TPROF_ON: + note_tprof("take_s", time.perf_counter() - _pt0) + if not flat: + self.last_ready[layer] = set() # 本层无就绪:miss_A 全归"时序(pread 未完成)" + if TPROF_ON: + note_tprof("promote_s", time.perf_counter() - _pt0, count_key="promote_n") + return 0 + gen = int(flat[0]) + # 记录本层"已就绪(pread 完成)"专家集:无论 buffer 是否匹配，pread 都已完成。 + self.last_ready[layer] = {int(flat[i]) for i in range(1, len(flat), 2)} + stg = self._gen_buf.pop(gen, None) # 按 handler 记的 gen 取回**正是它写过**的那块 buffer + if stg is None: + if TPROF_ON: + note_tprof("promote_s", time.perf_counter() - _pt0, count_key="promote_n") + return 0 # buffer 已被回收/不匹配 → 跳过（不会错配） + store._resident._ensure_layer(layer) # promote 在 acquire 前，池可能还没建 + resident = (store.resident_experts(layer) + if hasattr(store, "resident_experts") else set()) + pairs = [(int(flat[i]), int(flat[i + 1])) for i in range(1, len(flat), 2)] + # GPU 重映射路径无现成 used：仅对预读候选做 GPU membership 现算（不物化全量路由）。 + if used is None and route_inds is not None and num_experts is not None: + _tr0 = time.perf_counter() if TPROF_ON else 0.0 + used = route_used_subset([e for e, _ in pairs], route_inds, int(num_experts)) + if TPROF_ON: + note_tprof("route_s", time.perf_counter() - _tr0) + # 受保护（不被驱逐）= 本层真实路由（若已知），否则退化为全部预读专家。 + protect = set(used) if used is not None else {e for e, _ in pairs} + # 惰性切片直接放入池：正确性由 gen-匹配（切到 handler 真正写的那块 buffer）+ per-layer 环 + # （该 buffer 在该层 ring 次 submit 内不被覆盖，而 pool scatter 会在本前向读池时 eval）保证。 + # 不在此 mx.eval（那是每层 host 同步，在 verify 大流水线上会显著拖慢热路径）。 + # used 已知时只保留命中真实路由的专家：假阳性不写池（省 scatter、不挤掉有用专家）。 + _tpl0 = time.perf_counter() if TPROF_ON else 0.0 + pend = [(e, self._slice(stg[row])) for e, row in pairs + if e not in resident and (used is None or e in used)] + placed = 0 + for e, d in pend: + try: + store._resident._place_expert(layer, e, d, current=protect) + except ValueError: + break + if _STG_VERIFY: + self.placed_gen[(layer, e)] = gen # 诊断（默认关）：记录该专家这次落池用的 gen + placed += 1 + if TPROF_ON: + # place_s 含惰性切片构建 + 池 scatter 入图(不含其 GPU 执行);place_experts=实际写入专家数。 + note_tprof("place_s", time.perf_counter() - _tpl0, + count_key="place_experts", count=placed) + note_tprof("promote_s", time.perf_counter() - _pt0, count_key="promote_n") + self.promoted += placed + return placed + + # ---- zero-copy 双源解码：池侧区直接散写（C++ 维护侧区 e→物理行） ---- + def submit_pool_sideregion(self, layer, pred, resident, pool_list, base_row, gen=0): + """zero-copy 预取：触发 C++ 把预读专家段散写进池侧区行（填代 gen）。pool_list 为 _segs 顺序的 per-key 池数组。 + seg_nbytes 取自 blob 段表（与 pool key 顺序一致）。spec_slots = 侧区行数（=self.budget）。""" + import mlx_streaming.native_moe_ext as _N + layer = int(layer) + seg_nbytes = [int(nb) for *_, nb in self.src._segs] # 段字节（与 pool key 顺序一致） + path = f"{self.src.dir}/layer{layer:02d}.blob" + self.submitted += 1 + res = [int(e) for e in (resident or [])] + return _N.prefetch_pool_sideregion( + pool_list, seg_nbytes, pred, layer, path, self.stride, res, + int(self.budget), int(base_row), gen=int(gen)) + + def sideregion_contents(self, layer, gen=0): + """读该层某代 C++ 侧区缓存当前内容 {expert: 物理侧区行}（纯锁，不消费）。""" + import mlx_streaming.native_moe_ext as _N + flat = _N.sideregion_contents(int(layer), int(gen)) + return {int(flat[i]): int(flat[i + 1]) for i in range(0, len(flat), 2)} + + def sideregion_kv(self, layer, gen=0): + """读该层某代侧区为 (keys uint32, vals int32) 两个 device mx.array（C++ 直接建，供快路径合并）。""" + import mlx_streaming.native_moe_ext as _N + return _N.sideregion_kv(int(layer), int(gen)) + + def sideregion_reset(self): + """清空 C++ 侧区缓存（换 prompt/重置统计时调用）。""" + import mlx_streaming.native_moe_ext as _N + _N.sideregion_reset() + + def _slice(self, row: mx.array) -> dict: + # 按段 dtype 通用分派（支持 v1 affine 与 v2 mxfp4）： + # - uint32：weight，原样 reshape； + # - uint16：affine 的 scales/biases，view(bfloat16)； + # - uint8 ：mxfp4 的 scales，保持原样，绝不 view bf16。 + out = {} + off = 0 + for proj, tensor, dt, shape, nb in self.src._segs: + seg = row[off:off + nb] + out[f"{proj}.{tensor}"] = _typed_seg(seg, dt, shape) + off += nb + return out diff --git a/mlx_streaming/core/prefetch/patch.py b/mlx_streaming/core/prefetch/patch.py new file mode 100644 index 0000000..649d3bc --- /dev/null +++ b/mlx_streaming/core/prefetch/patch.py @@ -0,0 +1,56 @@ +"""模型 patch：把原生 MoE 块替换成流式块（文件后端 / 常驻切片版），并按需挂跨层预取。""" +from mlx_streaming import config +from mlx_streaming.core.moe.block import FileStreamingMoeBlock, StreamingMoeBlock +from mlx_streaming.core.prefetch.cross_layer import enable_cross_layer_prefetch + + +def patch_model_filebacked(model, store, hidden, moe_inter, group_size, bits, + proj_bits: dict | None = None, + layer_proj_bits: dict | None = None): + """把每个 MoE 块替换为 FileStreamingMoeBlock，并丢弃常驻的堆叠 switch_mlp。 + + store：FileExpertStore（所有 MoE 层共用，按 (layer,expert) 缓存）。 + proj_bits：非空时走混合精度（逐 proj 不同 bit），专家须为对应混合重量化产出。 + layer_proj_bits：{绝对层号: {proj:bits}}，非空时逐层用各自 proj_bits（优先于 proj_bits）， + 对应 requantize_dir_layered 产出。各层 QSL 用该层 bit，与流式存盘文件一一对应。 + 返回被替换的层数。被替换后原 switch_mlp 不再被引用，惰性权重不会被物化。 + """ + patched = 0 + for i, layer in enumerate(model.layers): + layer._layer_idx = i + # 用 object.__setattr__ 存反向 model 引用,避免进 nn.Module(=dict)的子模块树: + # 否则 model→layer→model 成环,wired_limit 等 tree_reduce 遍历会无限递归(RecursionError)。 + object.__setattr__(layer, "_prefetch_model_ref", model) + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): + pb = layer_proj_bits.get(i, proj_bits) if layer_proj_bits else proj_bits + # 捕获共享专家引用使其常驻（Qwen3-Next 有，Qwen3-MoE 无） + layer.mlp = FileStreamingMoeBlock( + gate=mlp.gate, top_k=mlp.top_k, norm_topk_prob=mlp.norm_topk_prob, + store=store, layer_idx=i, hidden=hidden, moe_inter=moe_inter, + group_size=group_size, bits=bits, proj_bits=pb, + shared_expert=getattr(mlp, "shared_expert", None), + shared_expert_gate=getattr(mlp, "shared_expert_gate", None), + ) + # 同理:不进 dict,避免 mlp→model→...→mlp 成环。供 native-fused-prefetch 取下层 gate。 + object.__setattr__(layer.mlp, "_prefetch_model_ref", model) + patched += 1 + if config.cross_layer_prefetch() or getattr(store, "_staging", None) is not None: + enable_cross_layer_prefetch() + return patched + + +def patch_model(model, store_factory=None): + """把模型里每个 MoE 块（含 switch_mlp 与 gate 的块）替换成 StreamingMoeBlock。 + + store_factory(layer_idx)->LruExpertStore，可为 None（仅做 uniq 切片、不接磁盘后端）。 + 返回被替换的层数，便于校验确实命中了 MoE 层。 + """ + patched = 0 + for i, layer in enumerate(model.layers): + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): + store = store_factory(i) if store_factory is not None else None + layer.mlp = StreamingMoeBlock(mlp, layer_idx=i, store=store) + patched += 1 + return patched diff --git a/mlx_streaming/core/profiling.py b/mlx_streaming/core/profiling.py new file mode 100644 index 0000000..5697049 --- /dev/null +++ b/mlx_streaming/core/profiling.py @@ -0,0 +1,131 @@ +"""热路径计时与诊断埋点（从 streaming_moe 抽出，集中管理）。 + +- PROF：细粒度分段计时（STREAM_PROF=1），probe 读取。 +- WINDOW_PROF：同层 submit→promote 的 attention/GDN 窗口时间（WINDOW_PROF=1）。 +- PREDICT_RECALL_PROF：运行时预测对实际路由的覆盖率（分诊预取覆盖问题，PREDICT_RECALL_PROF=1）。 +这些 dict 为共享可变对象：各模块 import 后原地累加，外部（run_mtp_spec）读同一对象。 +""" +import time + +import mlx.core as mx + +from mlx_streaming import config + +_PROF_ON = config.stream_prof() +PROF = {"route": 0.0, "pyremap": 0.0, "fetch": 0.0, "matmul": 0.0, "combine": 0.0, + "n_calls": 0} +WINDOW_PROF = {"sum_s": 0.0, "n": 0} +PREDICT_RECALL_PROF = {"hit": 0, "routed": 0, "n": 0} +# 主动预取路径 host 墙钟探针（PREFETCH_TPROF=1）：量主线程"不可与 GPU 重叠"的 CPU 时间。 +# 注意：MLX 惰性——gate matmul / pool scatter 的 GPU 计算落在前向末尾统一 eval，本探针看不到， +# 那部分用消融吞吐口径(NATIVE_NO_SUBMIT/NATIVE_NO_PROMOTE 差值)测。本探针专测 host 段： +# - predict_s：block._native_fused_prefetch 里建预测图(gate(x)+argpartition[+predicted_set tolist]) +# - submit_s ：stg.submit(prefetch_into_staging 的 host 调用，注册 GPU 完成回调) +# - promote_s：promote 总时长，再细分： +# · take_s ：prefetch_staging_take（C++ 锁读已就绪记录） +# · route_s：route_used_subset 的成员测试 .tolist()（GPU→host 同步，drain≤budget） +# · place_s：_place_expert 循环（建惰性切片 + 池 scatter 入图，不含其 GPU 执行） +# *_n 为各段调用次数（按"层·前向"计）。关时各调用点直接跳过，零开销。 +PREFETCH_TPROF = { + "predict_s": 0.0, "predict_n": 0, + "submit_s": 0.0, "submit_n": 0, + "promote_s": 0.0, "promote_n": 0, + "take_s": 0.0, "route_s": 0.0, "place_s": 0.0, "place_experts": 0, +} +TPROF_ON = config.prefetch_tprof() + + +def note_tprof(seg: str, dt: float, *, count_key: "str | None" = None, count: int = 1) -> None: + """累加预取某段 host 墙钟。dt 为秒;可选 count_key 同步累加调用/专家计数。""" + PREFETCH_TPROF[seg] += dt + if count_key is not None: + PREFETCH_TPROF[count_key] += count + + +def tprof_reset() -> None: + for k in PREFETCH_TPROF: + PREFETCH_TPROF[k] = 0 if isinstance(PREFETCH_TPROF[k], int) else 0.0 +# miss 归因（MISS_ATTRIB=1，promote 之后 / acquire 之前统计本层真实路由的命中构成）： +# - resident_hit：acquire 前已驻留（LRU 历史 + 本次 promote 命中） +# - miss_A_predicted：miss 且"在预测集里"——预测对了但没进池（budget 丢/时序没到/被驱逐） +# - miss_B_unpredicted：miss 且"不在预测集里"——预测器召回缺口（根本没预测到） +# dec_* 为 decode/verify 热路径专用桶（seq 短，与 prefill 长 seq 分离）：原 miss_attrib +# 仅埋在 host 分支、被 prefill 样本主导，无法反映 decode 热路径的真实 A/B 构成。 +MISS_ATTRIB = {"routed": 0, "resident_hit": 0, "miss_A_predicted": 0, + "miss_B_unpredicted": 0, "n": 0, + "dec_routed": 0, "dec_resident_hit": 0, "dec_miss_A": 0, + "dec_miss_B": 0, "dec_n": 0, + # miss_A 细分（decode 桶）：时序=就绪集里没有它（pread 没完成）； + # 驱逐=就绪过（pread 完成）但 acquire 前已不在池（被驱逐/没写进）。 + "dec_miss_A_timing": 0, "dec_miss_A_evicted": 0} + + +def note_miss_attrib(uniq, predicted, resident, is_decode: bool, ready=None) -> None: + """统计本层真实路由 uniq 的命中构成，分 prefill / decode 两桶累加。 + + uniq：本层真实路由专家集合；predicted：本层预测集；resident：acquire 前驻留集； + ready：本层最近 promote take 到的"已就绪(pread 完成)"专家集（用于细分 miss_A）。 + 供 host 路径与 GPU remap 路径共用，确保两条路径口径一致。 + """ + pset = predicted or set() + rset = ready or set() + for e in uniq: + MISS_ATTRIB["routed"] += 1 + if is_decode: + MISS_ATTRIB["dec_routed"] += 1 + if e in resident: + MISS_ATTRIB["resident_hit"] += 1 + if is_decode: + MISS_ATTRIB["dec_resident_hit"] += 1 + elif e in pset: + MISS_ATTRIB["miss_A_predicted"] += 1 + if is_decode: + MISS_ATTRIB["dec_miss_A"] += 1 + # 就绪集里有它=pread 完成过却仍 miss → 驱逐；否则 pread 没完成 → 时序。 + if e in rset: + MISS_ATTRIB["dec_miss_A_evicted"] += 1 + else: + MISS_ATTRIB["dec_miss_A_timing"] += 1 + else: + MISS_ATTRIB["miss_B_unpredicted"] += 1 + if is_decode: + MISS_ATTRIB["dec_miss_B"] += 1 + MISS_ATTRIB["n"] += 1 + if is_decode: + MISS_ATTRIB["dec_n"] += 1 + + +# 并集专家数探针(UNION_PROF=1):按本次前向的 seq 长度分桶,记每层路由专家"去重并集"大小。 +# seq=K(MTP verify,如 K=3)的桶即"K 个 token 的专家并集",决定池 cap 下限;seq=1 为 decode、 +# seq=chunk 为分块 prefill。值为 {seq: [sum_union, n_layer_calls]}。默认关、零开销。 +UNION_PROF: "dict[int, list]" = {} +# 分层原始样本(UNION_PROF=1)：{seq: {layer_idx: [本层每次前向的并集大小...]}}。 +# 仅聚合 [sum,n] 无法给出 cap 下限所需的 U_max/p99/分层分布，故额外留每层每次前向的原始值。 +# 样本量极小(≤层数×token数)，仅采集期占用，退出前汇总。 +UNION_SAMPLES: "dict[int, dict[int, list]]" = {} +UNION_ON = config.union_prof() + + +def note_union(seq: int, union_count: int, layer_idx: int = -1) -> None: + seq = int(seq) + union_count = int(union_count) + e = UNION_PROF.setdefault(seq, [0, 0]) + e[0] += union_count + e[1] += 1 + if layer_idx >= 0: + UNION_SAMPLES.setdefault(seq, {}).setdefault(int(layer_idx), []).append(union_count) + + +def union_reset() -> None: + UNION_PROF.clear() + UNION_SAMPLES.clear() + + +def prof_reset(): + for k in PROF: + PROF[k] = 0.0 + + +def _tick(seg, t0): + mx.eval # noqa (占位,避免误用) + PROF[seg] += time.perf_counter() - t0 diff --git a/mlx_streaming/core/route_trace.py b/mlx_streaming/core/route_trace.py new file mode 100644 index 0000000..b4249eb --- /dev/null +++ b/mlx_streaming/core/route_trace.py @@ -0,0 +1,50 @@ +"""MoE 路由 trace 采集工具（仅 probe 使用）。 + +热路径默认完全关闭；开启后会把每次 MoE block 的 layer 和专家集合记录下来， +用于离线替换策略模拟。注：开启会触发路由 ids 的 CPU 同步，不用于正式性能测试。 +""" +import json + + +_enabled = False +_events: list[dict] = [] + + +def enable() -> None: + global _enabled, _events + _enabled = True + _events = [] + + +def disable() -> None: + global _enabled + _enabled = False + + +def record(layer: int, experts, miss=None, resident=None, resident_rank=None) -> None: + if not _enabled: + return + vals = [int(e) for e in experts] + rec = {"layer": int(layer), "experts": sorted(set(vals))} + if miss is not None: + rec["miss"] = sorted({int(e) for e in miss}) + if resident is not None: + rec["resident"] = sorted({int(e) for e in resident}) + if resident_rank is not None: + rec["resident_rank"] = [[int(e), float(v)] for e, v in resident_rank.items()] + _events.append(rec) + + +def dump_jsonl(path: str) -> int: + with open(path, "w") as f: + for rec in _events: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + return len(_events) + + +def size() -> int: + return len(_events) + + +def events() -> list[dict]: + return list(_events) diff --git a/mlx_streaming/model_builder.py b/mlx_streaming/model_builder.py new file mode 100644 index 0000000..496444c --- /dev/null +++ b/mlx_streaming/model_builder.py @@ -0,0 +1,252 @@ +"""共享装配层:把"文件后端流式加载主模型 + 取 hidden + 贪心生成"等跨入口复用的逻辑 +集中在这里,而不是藏在某个入口脚本(原 validate_mtp)里。 + +依赖关系:本模块位于 core(mem/expert_store/streaming_moe)与 mtp 之上, +是把它们粘合成可运行模型的装配层,被 cli/ 各入口与测试复用。 + +环境变量: + MODEL 主模型路径(MLX 量化) + EXPERT_DIR 拆分/重量化后的 per-expert safetensors 目录 + EXPERT_SLOTS 每层常驻池容量 + EXPERT_POOL_PROFILE 每层池预算 JSON(无损省内存,可选) + HIDDEN_VARIANT pre_final_norm(默认)| post_final_norm(排错时切换) +""" +import json +import os + +import mlx.core as mx +from mlx_lm import load +from mlx_lm.models.base import create_attention_mask, create_ssm_mask + +from mlx_streaming import config +from mlx_streaming.core.cache.expert_store import FileExpertStore +from mlx_streaming.core.prefetch.patch import patch_model_filebacked + +MODEL = config.model_path() +EXPERT_DIR = config.expert_dir() +EXPERT_SLOTS = config.expert_slots() +# pre_final_norm(默认)| post_final_norm(排错时切换) +HIDDEN_VARIANT = config.hidden_variant() + +# 默认 profile 文件名:放在 EXPERT_DIR 下随专家目录一起走,常用跑法自动启用(无损省内存) +DEFAULT_PROFILE_NAME = "pool_profile.json" + + +def load_pool_profile(expert_dir: str) -> "dict[int, int] | None": + """解析每层池预算 profile,返回 layer_caps 或 None。 + + 优先级:环境变量 EXPERT_POOL_PROFILE 显式指定路径 > {expert_dir}/pool_profile.json 默认。 + EXPERT_POOL_PROFILE=none/0/off 显式关闭(回到 uniform capacity)。 + profile 无损:仅按各层真实工作集分配,命中率/输出/吞吐不变(caps 仍被 capacity 上限钳制)。 + """ + p = config.expert_pool_profile() + if p.lower() in ("none", "0", "off"): + return None + if not p: # 未显式指定 → 默认找专家目录下的 profile + cand = os.path.join(expert_dir, DEFAULT_PROFILE_NAME) + p = cand if os.path.exists(cand) else "" + if p and os.path.exists(p): + with open(p) as f: + caps = json.load(f).get("layer_caps", {}) + return {int(k): int(v) for k, v in caps.items()} + return None + + +def build_streaming_model(): + """用文件后端流式 patch 加载主模型(32GB 机器装不下 41GB 非流式)。""" + model, tok = load(MODEL, lazy=True) + # 取首个 MoE 维度 + dims = None + for layer in model.layers: + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): + gp = mlp.switch_mlp.gate_proj + dims = {"hidden": gp.input_dims, "moe_inter": gp.output_dims, + "group_size": getattr(gp, "group_size", 64), + "bits": getattr(gp, "bits", 4)} + break + bits, group, proj_bits, layer_proj_bits = ( + dims["bits"], dims["group_size"], None, None) + meta_path = os.path.join(EXPERT_DIR, "_split_meta.json") + if os.path.exists(meta_path): + with open(meta_path) as f: + meta = json.load(f) + ed = meta.get("dims", {}) + bits = ed.get("bits", bits) + group = ed.get("group_size", group) + proj_bits = ed.get("proj_bits") + if "per_layer_proj_bits" in ed: + layer_proj_bits = {int(k): v for k, v in ed["per_layer_proj_bits"].items()} + # 每层池预算 profile(pool_footprint 产出):默认从 {EXPERT_DIR}/pool_profile.json 自动启用, + # 无损省内存(命中率/输出/吞吐不变，仅不再为低占用层预留满 capacity)。 + layer_caps = load_pool_profile(EXPERT_DIR) + store = FileExpertStore(EXPERT_DIR, capacity=EXPERT_SLOTS, layer_caps=layer_caps) + if config.zerocopy_dual_source(): + # 零拷贝双源双缓冲：常驻池换成侧区模式（预分配 cap+2*spec_slots 行、禁 grow），复用原池 loader/cap/profile。 + from mlx_streaming.core.cache.resident_pool import ResidentExpertPool + _old = store._resident + # 默认单缓冲(持久 LFU,一份工作集,省一半侧区内存=生产路径);仅显式 legacy(SIDEREGION_LFU=0)用双缓冲。 + _spec_gens = 1 if config.sideregion_lfu() else 2 + store._resident = ResidentExpertPool( + _old.capacity, loader=_old.loader, layer_caps=_old.layer_caps, + spec_slots=config.pool_spec_slots(), + spec_gens=_spec_gens) + if config.stream_blob_loader(): + # blob 接入常驻池 miss-loader：复用 GPU-remap 快路径，小 EXPERT_SLOTS 即低内存。 + store._blob_loader = _make_blob_source(dims, group, bits) + # 主动预取（native-fused-prefetch miss→hit）：opt-in（NATIVE_FUSED_PREFETCH=1）。 + # 经"promote 只写真实路由命中专家"修正后已是净正：易缓存基座上 +15.5% tok/s + # （demand 11.86→13.70，hit 0.731→0.851，读盘 −45%；见 active-prefetch-turnaround-2026-06-17.md）。 + # 默认关只因收益依赖场景（基座可缓存性/是否磁盘受限）且只影响速度不影响质量、多占少量 staging 内存， + # 故作 opt-in 而非默认路径，落地配方见上述报告 §6。 + if config.native_fused_prefetch() and getattr(store, "_blob_loader", None) is not None: + try: + import mlx_streaming.native_moe_ext # noqa: F401 确认扩展已编译 + from mlx_streaming.core.prefetch.native_staging import NativeStagingManager + _budget = (config.pool_spec_slots() if config.zerocopy_dual_source() + else config.stream_blob_bg_budget(default=16)) + store._staging = NativeStagingManager(store._blob_loader, budget=_budget) + except Exception: + store._staging = None # 扩展不可用 → 关闭，不影响主路径 + # 零拷贝双源不变量：staging 侧区行数必须等于池 spec_slots，否则 C++ 会越界写池（静默损坏）。 + if config.zerocopy_dual_source() and getattr(store, "_staging", None) is not None: + assert store._staging.budget == store._resident.spec_slots, ( + f"零拷贝双源要求 staging.budget({store._staging.budget}) " + f"== 池 spec_slots({store._resident.spec_slots})") + if config.stream_blob_bg(): + # 后台预取池预填：bg 在独立 stream 物化预测专家，promote 写进池槽（需 CROSS_LAYER_PREFETCH=1）。 + from mlx_streaming.core.prefetch.bg_prefetch import BackgroundExpertPrefetcher + src = _make_blob_source(dims, group, bits) + store._blob_loader = src + store._bg = BackgroundExpertPrefetcher( + src, window=config.stream_blob_window()) + patch_model_filebacked(model, store, dims["hidden"], dims["moe_inter"], + group, bits, proj_bits=proj_bits, + layer_proj_bits=layer_proj_bits) + # 双源双缓冲：构造一个共享 VirtualPool（gen 跨层全局、每前向 +1）挂到每个流式 MoE 块。 + if config.zerocopy_dual_source() and getattr(store, "_staging", None) is not None: + from mlx_streaming.core.cache.virtual_pool import VirtualPool + from mlx_streaming.core.moe.block import FileStreamingMoeBlock + # 双源模式仍需 ahead 调度：block._native_fused_prefetch 靠 target_for 选目标层， + # 不传调度参数会让 target_for 恒返回 0（_num_layers=0）→ 预取全跳过、侧区永远空。 + _vpool = VirtualPool(store._resident, store._staging, config.pool_spec_slots(), + num_layers=len(model.layers), + cutoff=config.cross_layer_cutoff(), + ahead_lo=config.cross_layer_ahead_lo(), + ahead_hi=config.cross_layer_ahead_hi(), + store=store) + for layer in model.layers: + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, FileStreamingMoeBlock): + mlp._vpool = _vpool + # 主动预取（非 zerocopy）：挂 per-layer ahead 调度器 vpool（cutoff），让晚层预读更早发起。 + if (config.native_fused_prefetch() and not config.zerocopy_dual_source() + and getattr(store, "_staging", None) is not None): + from mlx_streaming.core.cache.virtual_pool import VirtualPool + from mlx_streaming.core.moe.block import FileStreamingMoeBlock + _sched = VirtualPool(num_layers=len(model.layers), + cutoff=config.cross_layer_cutoff(), + ahead_lo=config.cross_layer_ahead_lo(), + ahead_hi=config.cross_layer_ahead_hi()) + for layer in model.layers: + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, FileStreamingMoeBlock): + mlp._vpool = _sched + if config.stream_blob(): + _attach_blob_source(model, dims, group, bits) + # KV 量化(IsoQuant K4/V3 + SO(4) 旋转):仅作用于 12 个全注意力层,128k KV 3.0→~0.68 GiB。 + if config.kv_quant(): + from mlx_streaming.core.cache.kv_quant_patch import patch_kv_quant + patch_kv_quant(model, + group_size=config.kv_group_size(), + k_bits=config.kv_k_bits(), + v_bits=config.kv_v_bits(), + rotate=config.kv_rotate(), + seed=config.kv_rot_seed()) + # AUTOPIN 预热钉死(池 + blob_loader + staging 就绪后、return 前):按历史路由热度把每层 + # top-N 热专家预填进常驻池并 pin 住(不参与任何驱逐),消灭冷启动慢热;usage 缺失则跳过 + # (会话中计数仍在累计,落盘供下次启动)。AUTOPIN=0 零行为。 + if config.autopin(): + from mlx_streaming.core.cache import autopin as _autopin + _ap = _autopin.warm_start_pin(store) + if _ap["pinned"]: + print(f"[AUTOPIN] 预热 pin: {_ap['layers']} 层共 {_ap['pinned']} 专家 " + f"(frac={config.autopin_budget_frac()}, 耗时 {_ap['seconds']}s)" + + (f", 跳过 {_ap['skipped_layers']} 层" if _ap["skipped_layers"] else ""), + flush=True) + # 把 store 挂到 model 上，供只拿到 model 的调用方(server /api/stats 读专家池健康度)取用。 + # 非 array/dict/list/tuple 走 object 属性，不进 nn.Module 参数树。 + model._expert_store = store + return model, tok, store + + +def _make_blob_source(dims, group, bits): + from mlx_streaming.core.cache.blob_loader import BlobExpertSource + blob_dir = config.blob_dir() or os.path.join(EXPERT_DIR, "blobs") + workers = config.stream_blob_workers() + nocache = config.stream_blob_nocache(default="0") + num_experts = 512 + idx_path = os.path.join(blob_dir, "blob_index.json") + if os.path.exists(idx_path): + with open(idx_path) as f: + num_experts = int(json.load(f).get("num_experts", num_experts)) + return BlobExpertSource(blob_dir, dims["hidden"], dims["moe_inter"], group, bits, + num_experts=num_experts, workers=workers, nocache=nocache) + + +def _attach_blob_source(model, dims, group, bits): + """STREAM_BLOB=1：给每个流式 MoE 块注入共享 BlobExpertSource（全流式低内存路径）。""" + from mlx_streaming.core.cache.blob_loader import BlobExpertSource + from mlx_streaming.core.moe.block import FileStreamingMoeBlock + + blob_dir = config.blob_dir() or os.path.join(EXPERT_DIR, "blobs") + workers = config.stream_blob_workers() + nocache = config.stream_blob_nocache(default="1") + num_experts = 512 + idx_path = os.path.join(blob_dir, "blob_index.json") + if os.path.exists(idx_path): + with open(idx_path) as f: + num_experts = int(json.load(f).get("num_experts", num_experts)) + src = BlobExpertSource(blob_dir, dims["hidden"], dims["moe_inter"], group, bits, + num_experts=num_experts, workers=workers, nocache=nocache) + for layer in model.layers: + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, FileStreamingMoeBlock): + mlp._blob = src + + +def capture_prenorm_hidden(model, input_ids: mx.array) -> mx.array: + """跑主模型层循环但跳过最后的 model.norm,返回 last-layer hidden(norm 前)。 + + HIDDEN_VARIANT=post_final_norm 时返回 norm 之后(用于消歧排错)。 + """ + inner = model.model + h = inner.embed_tokens(input_ids) + layers = inner.layers + if not layers: + return h + cache = model.make_cache() + fa_idx = next((i for i, l in enumerate(layers) if not l.is_linear), 0) + ssm_idx = next((i for i, l in enumerate(layers) if l.is_linear), 0) + fa_mask = create_attention_mask(h, cache[fa_idx]) + ssm_mask = create_ssm_mask(h, cache[ssm_idx]) + for layer, c in zip(layers, cache): + mask = ssm_mask if layer.is_linear else fa_mask + h = layer(h, mask=mask, cache=c) + if HIDDEN_VARIANT == "post_final_norm": + h = inner.norm(h) + return h + + +def greedy(model, input_ids: mx.array, n: int) -> mx.array: + """主模型贪心生成 n 个 token,返回拼接后的完整序列(用作自投机参考)。""" + cache = model.make_cache() + cur = input_ids + out = [] + for _ in range(n): + logits = model(cur, cache=cache) + nxt = mx.argmax(logits[:, -1, :], axis=-1, keepdims=True) + out.append(nxt) + cur = nxt + mx.eval(nxt) + return mx.concatenate([input_ids] + out, axis=1) diff --git a/mlx_streaming/models/__init__.py b/mlx_streaming/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mlx_streaming/mtp/__init__.py b/mlx_streaming/mtp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mlx_streaming/mtp/bench_verdict.py b/mlx_streaming/mtp/bench_verdict.py new file mode 100644 index 0000000..f6b918a --- /dev/null +++ b/mlx_streaming/mtp/bench_verdict.py @@ -0,0 +1,28 @@ +"""最小树 A/B 评测的纯裁决逻辑:中位数 + go/no-go 判定(无副作用,可单测)。""" + + +def median(xs): + """样本中位数(偶数个取中间两数均值)。xs 非空。""" + s = sorted(xs) + n = len(s) + mid = n // 2 + if n % 2 == 1: + return s[mid] + return (s[mid - 1] + s[mid]) / 2 + + +def verdict_from_delta(delta, exact_all, margin=0.05): + """按跨 prompt 中位相对提速 delta 与 lossless 门给出裁决。 + + - exact_all=False → "bug"(lossless 硬门一票否决) + - delta > margin → "go" + - delta < -margin → "no-go" + - 其余(含边界) → "even" + """ + if not exact_all: + return "bug" + if delta > margin: + return "go" + if delta < -margin: + return "no-go" + return "even" diff --git a/mlx_streaming/mtp/drafter.py b/mlx_streaming/mtp/drafter.py new file mode 100644 index 0000000..19ce82d --- /dev/null +++ b/mlx_streaming/mtp/drafter.py @@ -0,0 +1,188 @@ +"""真实 MTP drafter：把 Qwen3NextMTP 包成 mtp_generate 需要的 draft/sync 接口。""" +import mlx.core as mx + +from mlx_streaming.mtp.qwen3_next_mtp import mtp_step +from mlx_streaming.mtp.kv_cache import _snapshot, _restore + + +class MTPDrafter: + """把 Qwen3NextMTP 包成 mtp_generate 需要的 drafter 接口。""" + + def __init__(self, mtp, lm_head): + self.mtp = mtp + self.lm_head = lm_head + self.embed_tokens = mtp.embed_tokens + + def make_cache(self): + from mlx_lm.models import cache as kc + return [kc.KVCache()] # MTP 单层全注意力 + + def draft(self, H_last, x_ids, mtp_cache, K, topk: int = 0): + # 注意:每草稿一次 int(argmax) host 同步反而最快 —— 它让 MLX 把每步 draft 的 + # 图保持很小、及时释放;试过"全程 GPU argmax、末尾一次同步"攒大惰性图,draft + # 慢 5×、端到端 24.9→16.3(A/B 证伪,见报告)。保持逐步同步。 + # topk>0(探针用):额外返回每个位置 MTP 的 top-k 候选 id(降序),量树形展开的救回上界。 + drafts, cands = [], [] + h, cur = H_last, x_ids + for _ in range(K): + logits, mh = mtp_step(self.mtp, h, cur, self.lm_head, mtp_cache[0]) + lg = logits[0].reshape(-1) + if topk > 0: + order = [int(i) for i in mx.argsort(lg)[-topk:].tolist()][::-1] # 降序 top-k + d = order[0] + cands.append(order) + else: + d = int(mx.argmax(lg)) + drafts.append(d) + h, cur = mh, mx.array([[d]]) + if topk > 0: + return drafts, cands + return drafts + + def draft_adaptive(self, H_last, x_ids, mtp_cache, depth_max, tau): + """置信度门控动态深度:逐位贪婪抽,累计置信度 C=∏p_i 跌破 tau 即停,最多 depth_max 位。 + + 返回可变长度(1..depth_max)的草稿链。C 用每位 top-1 softmax 概率连乘;低置信提前收敛 + (本步只 verify 到当前深度,省后续位置的专家加载),高置信抽满到 depth_max。始终至少抽 1 位 + (depth=1 即退化为普通单 token 解码:verify 只喂 [x]、恒接受模型真值 1 个)。 + """ + drafts = [] + h, cur = H_last, x_ids + conf = 1.0 + for i in range(depth_max): + logits, mh = mtp_step(self.mtp, h, cur, self.lm_head, mtp_cache[0]) + lg = logits[0].reshape(-1) + d = int(mx.argmax(lg)) + drafts.append(d) + h, cur = mh, mx.array([[d]]) + if i + 1 < depth_max: # 末位无需再判是否加深 + conf *= float(mx.softmax(lg)[d]) + if conf < tau: + break + return drafts + + def draft_adaptive_tree(self, H_last, x_ids, mtp_cache, depth_max, tau): + """置信度门控变长 + pos0 分支:返回 (chainA, chainB)。 + + chainA 同 draft_adaptive(累计置信度门控,变长 n=1..depth_max);chainB 是第1 位 top-2 分支, + 续抽到与 chainA 同深度 n(仅 n>=2 时,浅步 depth=1 无位置可救 → chainB=None)。用快照隔离 + A/B 分叉,B 不带 A 的递归污染。供合并路径:动态定深 + 对深链做 pos0 救回。 + """ + logits1, mh1 = mtp_step(self.mtp, H_last, x_ids, self.lm_head, mtp_cache[0]) + lg = logits1[0].reshape(-1) + top2 = [int(i) for i in mx.argsort(lg)[-2:].tolist()][::-1] # [d1a, d1b] 降序 + d1a, d1b = top2[0], top2[1] + snap_after_x = _snapshot(mtp_cache) + + # chainA:从 d1a 续抽,累计置信度跌破 tau 即停(与 draft_adaptive 同语义)。 + chainA = [d1a] + conf = float(mx.softmax(lg)[d1a]) + h, cur = mh1, mx.array([[d1a]]) + while len(chainA) < depth_max and conf >= tau: + lo, mh = mtp_step(self.mtp, h, cur, self.lm_head, mtp_cache[0]) + lg2 = lo[0].reshape(-1) + d = int(mx.argmax(lg2)) + chainA.append(d) + conf *= float(mx.softmax(lg2)[d]) + h, cur = mh, mx.array([[d]]) + + n = len(chainA) + chainB = None + if n >= 2: # 浅步无后续位置可救,不抽 B + _restore(mtp_cache, snap_after_x) + chainB = [d1b] + h, cur = mh1, mx.array([[d1b]]) + for _ in range(n - 1): + lo, mh = mtp_step(self.mtp, h, cur, self.lm_head, mtp_cache[0]) + d = int(mx.argmax(lo[0])) + chainB.append(d) + h, cur = mh, mx.array([[d]]) + return chainA, chainB + + def draft_tree(self, H_last, x_ids, mtp_cache, K, pos1=True): + """最小树:第1(始终)、第2(pos1=True 时)草稿位置展开 top-2,返回三条链,各长 K。 + + - chainA = [d1a(top-1), d2a(top-1), d3a] 全 top-1,同 draft。 + - chainB = [d1b(top-2), d2b, d3b] 第1 草稿位置 top-2 分支 → pos0 救回。 + - chainC = [d1a, d2c(第2 位次选), d3c] 或 None 第2 草稿位置 top-2 分支 → pos1 救回。 + + 各分叉点用 mtp_cache 快照隔离,保证任一分支不带其它分支的递归污染:snap_after_x(x 处理后) + 分叉 d1a/d1b;snap_after_d1a(d1a 处理后)分叉 d2a/d2c。探针实测第2 位「首选错次选对」比例 + (~11%)高于第1 位(~7%),故 pos1 分支值得抽;但每步多一次 MTP 续抽,pos1=False 时退化为 + 仅 chainA/chainB(与旧最小树成本一致),便于 A/B 隔离 pos1 增量。 + """ + logits1, mh1 = mtp_step(self.mtp, H_last, x_ids, self.lm_head, mtp_cache[0]) + lg = logits1[0].reshape(-1) + top2 = [int(i) for i in mx.argsort(lg)[-2:].tolist()][::-1] # [d1a, d1b] 降序 + d1a, d1b = top2[0], top2[1] + snap_after_x = _snapshot(mtp_cache) # x 处理后、第1 位分叉前的 MTP 递归态 + + def _continue(first, h0, n): + """首 token first 已定,从隐状态 h0 再贪婪续抽 n 个,返回长度 1+n 的链。""" + chain = [first] + h, cur = h0, mx.array([[first]]) + for _ in range(n): + lo, mh = mtp_step(self.mtp, h, cur, self.lm_head, mtp_cache[0]) + d = int(mx.argmax(lo[0])) + chain.append(d) + h, cur = mh, mx.array([[d]]) + return chain + + chainC = None + if pos1: + # 从 d1a 续抽第2 位,捕获第2 位 top-2(d2a/d2c),供 chainA 与 chainC 分叉。 + lo2, mh_after_d1a = mtp_step(self.mtp, mh1, mx.array([[d1a]]), self.lm_head, mtp_cache[0]) + lg2 = lo2[0].reshape(-1) + top2_p1 = [int(i) for i in mx.argsort(lg2)[-2:].tolist()][::-1] # [d2a, d2c] 降序 + d2a, d2c = top2_p1[0], top2_p1[1] + snap_after_d1a = _snapshot(mtp_cache) # d1a 处理后、第2 位分叉前的递归态 + chainA = [d1a] + _continue(d2a, mh_after_d1a, K - 2) # [d1a, d2a, d3a...] + _restore(mtp_cache, snap_after_d1a) # 回到 d1a 态,C 链从第2 位次选分叉 + chainC = [d1a] + _continue(d2c, mh_after_d1a, K - 2) # [d1a, d2c, d3c...] + else: + chainA = _continue(d1a, mh1, K - 1) # 仅全 top-1 链(旧最小树成本) + + # chainB:回到 x 态,从第1 位次选 d1b 分叉续抽。 + _restore(mtp_cache, snap_after_x) + chainB = _continue(d1b, mh1, K - 1) # [d1b, d2b, d3b...] + return chainA, chainB, chainC + + def draft_paths(self, H_last, x_ids, mtp_cache, K, P): + """完整树 batch-of-paths:位置1 展开 top-P,返回 P 条链(各长 K)。 + + 位置1 取 MTP top-P 候选 [d1_0..d1_{P-1}];每个候选从 pos1 的共享 MTP 递归态 mh1 分叉, + 贪婪续抽 K-1 个 token 成一条链。用 mtp_cache 快照保证各链从同一起点分叉、互不污染。 + P=1 退化为普通链;P=2 等价 draft_tree。 + """ + logits1, mh1 = mtp_step(self.mtp, H_last, x_ids, self.lm_head, mtp_cache[0]) + lg = logits1[0].reshape(-1) + firsts = [int(i) for i in mx.argsort(lg)[-P:].tolist()][::-1] # top-P 降序 + snap_pos1 = _snapshot(mtp_cache) + + def _continue(first): + chain = [first] + h, cur = mh1, mx.array([[first]]) + for _ in range(K - 1): + lo, mh = mtp_step(self.mtp, h, cur, self.lm_head, mtp_cache[0]) + d = int(mx.argmax(lo[0])) + chain.append(d) + h, cur = mh, mx.array([[d]]) + return chain + + paths = [] + for j, f in enumerate(firsts): + if j > 0: + _restore(mtp_cache, snap_pos1) # 回到 pos1 态,从同一起点分叉 + paths.append(_continue(f)) + return paths + + def sync(self, prev_H, rH, replay_in, mtp_cache): + """用已接受 token 的真实主模型 hidden 推进 MTP KV cache。 + + MTP 在位置 i 消费 (H_i, t_{i+1}) 预测 t_{i+2};因此提交 accepted prefix + `[t_{i+1}, ..., t_{i+n}]` 时,hidden 序列应为 `[H_i, ..., H_{i+n-1}]`。 + """ + from mlx_streaming.mtp.qwen3_next_mtp import mtp_advance + h_seq = mx.concatenate([prev_H, rH[:, :-1, :]], axis=1) + H = mtp_advance(self.mtp, h_seq, replay_in, mtp_cache[0]) + mx.eval(H) diff --git a/mlx_streaming/mtp/generate.py b/mlx_streaming/mtp/generate.py new file mode 100644 index 0000000..c334057 --- /dev/null +++ b/mlx_streaming/mtp/generate.py @@ -0,0 +1,506 @@ +"""Qwen3-Next MTP 自投机贪婪解码循环。 + +见 spec/plan: docs/superpowers/{specs,plans}/2026-06-07-qwen3next-mtp-self-speculation*。 + +核心:每步主模型前向出 hidden -> MTP 自回归抽 K 草稿 -> 主模型并行验证 -> +接受最长命中前缀 -> cache 快照/恢复/重放回滚(统一处理 ArraysCache 与 KVCache), +保证输出与非投机贪婪逐 token 等价。 + +cache 校验/提交机制见 mtp/kv_cache.py；drafter 接口见 mtp/drafter.py。 +""" +import time + +import mlx.core as mx +from mlx_lm.models.base import create_attention_mask, create_ssm_mask + +from mlx_streaming import config +from mlx_streaming.mtp.kv_cache import ( + enable_qwen3next_speculative_checkpoints, begin_speculative_checkpoints, + _snapshot, _restore, commit_verified_prefix, commit_verified_snapshot, + tile_caches, commit_tree_row) + + +def _batch_direct_commit_guaranteed(model) -> bool: + """判断该模型的 batch verify 是否「必定」能直接提交(commit_verified_prefix 恒成功)。 + + commit 成功的充要条件是每个 cache 都满足:可裁剪(KVCache)或已捕获 per-token checkpoint。 + - 非线性层 → KVCache,可裁剪,必成功; + - 线性层 → 必须是被 patch 过的 Qwen3NextGatedDeltaNet,verify 前向才会写 `_spec_checkpoints`。 + 全部满足才返回 True——此时每步的 `snap_m` 回退快照永远用不上,可安全省略以压低在途峰值。 + 通用/玩具递归层(未 patch)返回 False,保留 snap_m 走安全 replay 兜底(与原行为逐 token 等价)。 + """ + import mlx_streaming.mtp.kv_cache as _kv + if not getattr(_kv, "_QWEN3NEXT_CHECKPOINTS_PATCHED", False): + return False + try: + from mlx_lm.models.qwen3_next import Qwen3NextGatedDeltaNet + layers = model.model.layers + except (ImportError, AttributeError): + return False + for l in layers: + if getattr(l, "is_linear", False) and not isinstance( + getattr(l, "linear_attn", None), Qwen3NextGatedDeltaNet): + return False + return True + + +def forward_with_hidden(model, ids, cache, compute_logits: bool = True): + """跑主模型层循环 + 最终 norm,返回 (logits(1,L,V), hidden(1,L,H))。 + + logits 恒由 final-norm 后的 H 经 lm_head 得到(保证与贪婪一致)。 + 返回给 MTP drafter 的 hidden:**默认用 final-norm 之前的 h**——这与 MTP 训练/验证 + (`capture_prenorm_hidden`)消费的输入一致;喂 norm 后的 hidden 会双重归一化、显著压低 + 草稿接受率(不影响正确性,因 verify 走主模型)。`MTP_HIDDEN=post_norm` 可切回旧行为对照。 + + compute_logits=False:跳过 lm_head(2048×151936 的大投影),只更新 cache 并返回 hidden。 + 供分块 prefill 的中间块用(中间块不需要 logits,只需 cache 因果累积)。 + """ + inner = model.model + h = inner.embed_tokens(ids) + layers = inner.layers + has_full = any(not l.is_linear for l in layers) + has_linear = any(l.is_linear for l in layers) + fa_idx = next((i for i, l in enumerate(layers) if not l.is_linear), 0) + ssm_idx = next((i for i, l in enumerate(layers) if l.is_linear), 0) + fa_mask = create_attention_mask(h, cache[fa_idx]) if has_full else None + ssm_mask = create_ssm_mask(h, cache[ssm_idx]) if has_linear else None + for layer, c in zip(layers, cache): + mask = ssm_mask if layer.is_linear else fa_mask + h = layer(h, mask=mask, cache=c) + H = inner.norm(h) + hidden = H if config.mtp_hidden() == "post_norm" else h + logits = model.lm_head(H) if compute_logits else None + return logits, hidden + + +def prefill_chunked(model, ids, cache, chunk: "int | None" = None): + """分块 prefill:把 ids 按 chunk 切片逐块喂入,增量更新 cache,返回最后一块的 (logits, hidden)。 + + 整段 prefill 一次前向的激活峰值 ∝ prompt 长度(每 MoE 层瞬时物化大量唯一专家 + 长序列激活)。 + 分块后每块只算 chunk 个 token,逐块 eval 释放上一块瞬时图,峰值压回 ∝chunk,与 decode 同稳态。 + + 数值等价:全注意力层逐块因果累积与整段完全等价;线性层(gated-delta)分块在 chunk 边界处有 + 重结合的末位浮点差异(与 decode 本身逐 token 同源,可忽略)。中间块跳过 lm_head 省大投影。 + chunk<=0 或 prompt 不超过一块时,退化为整段 prefill。 + """ + if chunk is None: + chunk = config.prefill_chunk() + L = ids.shape[1] + if chunk <= 0 or L <= chunk: + return forward_with_hidden(model, ids, cache) + last_start = ((L - 1) // chunk) * chunk + for s in range(0, last_start, chunk): + _, h = forward_with_hidden(model, ids[:, s:s + chunk], cache, compute_logits=False) + mx.eval(h) # 物化本块、写 cache、释放上一块瞬时图 + return forward_with_hidden(model, ids[:, last_start:], cache) + + +def forward_with_hidden_stepwise(model, ids, cache, capture_snapshots: bool = False): + """逐 token 走主模型解码路径,可在每个 token 后保存 cache 快照。 + + 这是 exact direct-commit 的保守路径:verify 与 baseline 都走 N=1 cache 更新, + 接受前缀时直接恢复到对应 token 后的快照,不再 replay accepted prefix。 + """ + logits_parts, hidden_parts, snaps = [], [], [] + for i in range(ids.shape[1]): + logits, H = forward_with_hidden(model, ids[:, i:i + 1], cache) + mx.eval(logits, H) + logits_parts.append(logits) + hidden_parts.append(H) + if capture_snapshots and i < ids.shape[1] - 1: + snaps.append(_snapshot(cache)) + logits = mx.concatenate(logits_parts, axis=1) + H = mx.concatenate(hidden_parts, axis=1) + if capture_snapshots: + return logits, H, snaps + return logits, H + + +# ----------------------------------------------------------------- 接受判定 +def accept_prefix(drafts, preds): + """drafts: MTP 抽的 K 个草稿; preds: 主模型对应位置真实下一 token(len==K)。 + + 返回命中草稿数 matched(最长使 preds[j]==drafts[j] 的前缀长度)。 + 发射 token 由调用方组装为 drafts[:matched] + [重放末位 bonus],统一处理纠正/全命中, + 保证贪婪下逐 token 与非投机等价。 + """ + matched = 0 + for d, p in zip(drafts, preds): + if int(p) == int(d): + matched += 1 + else: + break + return matched + + +# --------------------------------------------------------- 完整树形验证(batch-of-paths) +def tree_verify_step(model, drafter, x, H_last, x_ids, mtp_cache, main_cache, K, P, snap_d): + """一次 batched 前向并行验证 P 条候选路径,提交接受最长的赢家路径。 + + 返回 (new_tokens, x_new, rH, accepted_in, matched)。 + + 原理:每条路径 `[x, d_1..d_{K-1}]` 是普通线性序列,拍到 batch 维后线性层/全注意力层都走成熟 + batch 前向(逐 row 与 batch=1 等价,见 benchmarks/test_layerwise.py)。batch=P 的计算量加宽了 + 每层预取窗口(hit_rate),多路径又提升接受长度(accept_len)。验证后按赢家 row 提取 checkpoint + /trim 提交回 batch=1 主 cache,后续解码逐 token 从赢家路径续,与非投机贪婪等价。 + """ + paths = drafter.draft_paths(H_last, x_ids, mtp_cache, K, P) + verify_in = mx.array([[x] + p[: K - 1] for p in paths]) # (P, K) + + snap_m = _snapshot(main_cache) # tile 前的 batch=1 快照(回退用) + tile_caches(main_cache, len(paths)) + begin_speculative_checkpoints(main_cache) + vlogits, vH = forward_with_hidden(model, verify_in, main_cache) + mx.eval(vlogits, vH) + + best_w, matched, preds = 0, -1, None + for wi, p in enumerate(paths): + preds_i = [int(t) for t in mx.argmax(vlogits[wi], axis=-1)] + m = accept_prefix(p, preds_i) + if m > matched: # 平局取更靠前的路径(top-1 优先) + best_w, matched, preds = wi, m, preds_i + drafts = paths[best_w] + accepted_len = min(matched + 1, K) + + committed = commit_tree_row(main_cache, verified_len=K, + accepted_len=accepted_len, row=best_w) + if snap_d is not None: + _restore(mtp_cache, snap_d) + + if not committed: + # 理论上 checkpoint 齐全不会走到;稳妥回退:恢复 batch=1 主 cache 后重放赢家 accepted prefix。 + _restore(main_cache, snap_m) + accepted_in = mx.array([[x] + drafts[:matched]]) + rlogits, rH = forward_with_hidden(model, accepted_in, main_cache) + mx.eval(rlogits) + bonus = int(mx.argmax(rlogits[:, -1, :])) + return drafts[:matched] + [bonus], bonus, rH, accepted_in, matched + + rH = vH[best_w:best_w + 1, :accepted_len, :] + accepted_in = verify_in[best_w:best_w + 1, :accepted_len] + if matched == K: + return drafts[:K], drafts[-1], rH, accepted_in, matched + return drafts[:matched] + [preds[matched]], preds[matched], rH, accepted_in, matched + + +# ----------------------------------------------------------------- 主循环 +def mtp_generate(model, drafter, tok, prompt, max_tokens, K=3, ids_mode=False, + profile=False, on_tokens=None, main_cache=None, cached_len=0): + """贪婪 MTP 自投机。 + + drafter 需提供 draft(H_last(1,1,H), x_ids(1,1), mtp_cache, K) -> list[int](长度 K); + 可选 make_cache()->list 与 sync(rH, replay_in, mtp_cache)。 + ids_mode=True 时 prompt 已是 ids(1,L) 且返回 token id 列表(测试用)。 + on_tokens:可选流式钩子,每步用「本步真正写入 produced 的 token id 列表」回调一次 + (prefill 的首 token 作为第一次回调);返回 True 表示请求尽快停止生成。 + + 跨轮复用(main_cache/cached_len):传入已持有前 cached_len 个 prompt token 的 main_cache, + 则只 prefill `prompt[:, cached_len:]`(等价于一次多 token decode,复用历史 KV/递归态,不重算)。 + 调用方须保证 prompt[:cached_len] 与 cache 中已有 token 完全一致(否则结果错误)。 + 不变式:返回时 main_cache 恰好持有 `prompt + produced[:-1]`(produced[-1] 为 pending,未入 cache), + 调用方可据此拼下轮的 cached prefix。main_cache=None 时内部新建(默认,整段 prefill)。 + """ + enable_qwen3next_speculative_checkpoints() + if main_cache is None: + main_cache = model.make_cache() + cached_len = 0 + mtp_cache = drafter.make_cache() if hasattr(drafter, "make_cache") else None + ids = prompt if ids_mode else mx.array([tok.encode(prompt)]) + if cached_len < 0 or cached_len >= ids.shape[1]: + raise ValueError( + f"cached_len={cached_len} 必须落在 [0, prompt_len={ids.shape[1]}) 内,至少留 1 个 token 供 prefill") + + # prefill(分块):只喂尚未在 cache 中的后缀 ids[:, cached_len:](cached_len==0 即整段); + # 得到第 1 个 pending token x 与其 hidden;分块把激活峰值压到与 decode 同稳态。 + logits, H = prefill_chunked(model, ids[:, cached_len:], main_cache) + x = int(mx.argmax(logits[:, -1, :])) + H_last = H[:, -1:, :] + produced = [x] + _stop_from_initial = on_tokens is not None and on_tokens([x]) + mx.eval(H_last) + t0 = time.perf_counter() + n_steps = 0 + # profile 分段累加(秒):把 sync/finalize 拆开,避免隐藏在 replay 名下。 + t_draft = t_verify = t_replay = t_snap = 0.0 + t_commit = t_sync = t_finalize = 0.0 + direct_commits = fallback_replays = replayed_tokens = 0 + # matched 直方图:accept_hist[j] = 恰好命中 j 个草稿(j∈0..K)的步数。用于实测每位置接受率, + # 不做几何分布假设。emitted/step = min(matched+1,K)，故上限为 K(本实现 verify 只喂 K-1 草稿)。 + accept_hist = [0] * (max(K, config.depth_max() if config.adaptive_depth() else K) + 1) + # top-k 覆盖探针:每位置统计模型真实 token 落在 MTP top-1/2/3 的步数(树形救回上界)。 + topk_probe = config.accept_topk() if profile else 0 + tk_n = [0] * K + tk_cover1 = [0] * K + tk_cover2 = [0] * K + tk_cover3 = [0] * K + verify_mode = config.mtp_verify_mode() + tree_mode = config.tree_top2() + tree_p1_mode = tree_mode and config.tree_top2_p1() # 第2 位(pos1)救回,附加在 tree_top2 上 + tree_verify_mode = config.tree_verify() # 完整树形验证(batch-of-paths),优先级最高 + # 置信度门控动态深度:仅在 plain 单链路径生效(与 tree/tree_verify/step 互斥)。 + adaptive_mode = config.adaptive_depth() and not tree_mode and not tree_verify_mode \ + and verify_mode != "step" + adaptive_rescue_mode = adaptive_mode and config.adaptive_rescue() # 深链叠加 pos0 救回 + conf_tau_v = config.conf_tau() + depth_max_v = max(1, config.depth_max()) + tree_P = max(1, config.tree_branches()) + tree_rescues = 0 # 第1 位 top-2 成功救回(B 链首命中)的步数 + tree_rescues_p1 = 0 # 第2 位 top-2 成功救回(C 链第2 token 命中)的步数 + # 纯 batch 直接提交路径 + 模型保证 commit 恒成功时,跳过每步一次「全 cache 深拷贝 + eval」 + # (snap_m 只用于 tree 救回 / step / replay 回退,这些路径都不满足下方条件)。省 ~72MiB 在途 + # 峰值和一次同步栅栏,数值完全不变(bit-exact,只改内存调度)。模型结构恒定,循环外算一次即可。 + _skip_snap = (verify_mode != "step") and (not tree_mode) and (not tree_verify_mode) \ + and (not adaptive_rescue_mode) and _batch_direct_commit_guaranteed(model) + + while len(produced) < max_tokens and not _stop_from_initial: + x0 = x + x_ids = mx.array([[x]]) + + snap_d = _snapshot(mtp_cache) if mtp_cache else None + prev_H_last = H_last + + # 完整树形验证:一次 batched 前向验证 P 条路径,提交赢家。自带 step 尾处理后 continue。 + if tree_verify_mode and verify_mode != "step": + if profile: + _tic = time.perf_counter() + new_tokens, x, rH, accepted_in, matched = tree_verify_step( + model, drafter, x, H_last, x_ids, mtp_cache, main_cache, K, tree_P, snap_d) + accept_hist[matched] += 1 + direct_commits += 1 + if profile: + t_verify += time.perf_counter() - _tic + _tic = time.perf_counter() + _stop = False + _n_before = len(produced) + for t in new_tokens: + produced.append(t) + if len(produced) >= max_tokens: + break + # 只把「本步真正写入 produced」的 token 交给回调:多 token 步命中 max_tokens + # 上限时,截断掉的尾巴不应上报给流式消费者(避免超报被丢弃的 token)。 + if on_tokens is not None and on_tokens(produced[_n_before:]): + _stop = True + H_last = rH[:, -1:, :] + if mtp_cache is not None and hasattr(drafter, "sync"): + drafter.sync(prev_H_last, rH, accepted_in, mtp_cache) + if profile: + t_sync += time.perf_counter() - _tic + n_steps += 1 + mx.eval(x, H_last) + if _stop: + break + continue + + if profile: + _tic = time.perf_counter() + tree_b = None + tree_c = None + step_k = K # 本步验证深度(动态深度时可变) + if tree_mode and verify_mode != "step": + # chainA(全 top-1)、chainB(第1 位 top-2 → pos0 救回)、chainC(第2 位 top-2 → pos1 救回) + drafts, tree_b, tree_c = drafter.draft_tree( + H_last, x_ids, mtp_cache, K, pos1=tree_p1_mode) + draft_cands = None + elif adaptive_mode: + # 置信度门控:返回可变长度(1..depth_max)草稿链,本步深度随置信度自适应。 + if adaptive_rescue_mode: + # 合并路径:变长 chainA + pos0 分支 chainB(深链才有),供下方 pos0 救回。 + drafts, tree_b = drafter.draft_adaptive_tree( + H_last, x_ids, mtp_cache, depth_max_v, conf_tau_v) + else: + drafts = drafter.draft_adaptive(H_last, x_ids, mtp_cache, depth_max_v, conf_tau_v) + draft_cands = None + step_k = len(drafts) + elif topk_probe > 0: + drafts, draft_cands = drafter.draft(H_last, x_ids, mtp_cache, K, topk=topk_probe) + else: + drafts, draft_cands = drafter.draft(H_last, x_ids, mtp_cache, K), None # 长度 K + verify_in = mx.array([[x] + drafts[: step_k - 1]]) # [x, d_1..d_{step_k-1}] + if profile: + mx.eval(verify_in) + t_draft += time.perf_counter() - _tic + _tic = time.perf_counter() + + snap_m = None if _skip_snap else _snapshot(main_cache) + if profile: + t_snap += time.perf_counter() - _tic + _tic = time.perf_counter() + + verify_snaps = None + if verify_mode == "step": + vlogits, vH, verify_snaps = forward_with_hidden_stepwise( + model, verify_in, main_cache, capture_snapshots=True) + else: + begin_speculative_checkpoints(main_cache) + vlogits, vH = forward_with_hidden(model, verify_in, main_cache) + mx.eval(vlogits, vH) + preds = [int(t) for t in mx.argmax(vlogits[0], axis=-1)] # 长度 K + matched = accept_prefix(drafts, preds) + # 最小树救回:A 链首草稿被拒(matched==0)且 B 候选=模型真实 token(=preds[0])时, + # 恢复 cache、改验 B 链。preds[0] 只依赖 [x],两链一致;故 d1b==preds[0] 即 B 首命中。 + # lossless(2026-07-05 定论,见 reports/tree-top2-rescue-2026-07-05.md):曾疑此救回非 + # bit-lossless,实为 demand_dual 的一个 bug——inds(argpartition 切片)是非连续视图,C++ + # 按连续读 data() 致 seq≥2 的 token≥1 装错专家,spec/救回全体失真。修 demand_dual 用 + # contiguous() 物化后,seq=K verify 与 seq=1 解码逐位等价,spec 与本救回均 bit-lossless。 + if tree_b is not None and matched == 0 and tree_b[0] == preds[0]: + _restore(main_cache, snap_m) + begin_speculative_checkpoints(main_cache) + # 重建 verify_in 为 B 链:后续 accepted_in=verify_in[:, :accepted_len] 才会取到 chainB, + # 否则残留 chainA 的错误首草稿会经 sync 污染 MTP cache、拉低后续草稿质量(低估救回收益)。 + verify_in = mx.array([[x] + tree_b[:step_k - 1]]) + vlogits, vH = forward_with_hidden(model, verify_in, main_cache) + mx.eval(vlogits, vH) + preds = [int(t) for t in mx.argmax(vlogits[0], axis=-1)] + drafts = tree_b + matched = accept_prefix(drafts, preds) + tree_rescues += 1 + # 第2 位(pos1)救回:A 链首草稿命中(matched==1)但第2 草稿被拒,且 chainC 的第2 token + # =模型真值 preds[1] 时,改验 chainC=[x, d1a, d2c, ...]。preds[1] 只依赖 [x, d1a],A/C + # 前缀一致故 preds[0]/preds[1] 不变、重验后 matched≥2;d3c 可能续命中把接受长度顶到 3。 + # 与 pos0 救回同构(纯重验、只提交模型 argmax token → bit-lossless,并集不放大)。 + elif tree_c is not None and matched == 1 and len(tree_c) > 1 and tree_c[1] == preds[1]: + _restore(main_cache, snap_m) + begin_speculative_checkpoints(main_cache) + verify_in = mx.array([[x] + tree_c[:step_k - 1]]) + vlogits, vH = forward_with_hidden(model, verify_in, main_cache) + mx.eval(vlogits, vH) + preds = [int(t) for t in mx.argmax(vlogits[0], axis=-1)] + drafts = tree_c + matched = accept_prefix(drafts, preds) + tree_rescues_p1 += 1 + accept_hist[matched] += 1 + if draft_cands is not None: + # preds[i]=模型在位置 i 的真实下一 token;查它是否在 MTP 该位置 top-1/2/3 候选里。 + for i in range(K): + p = preds[i] + c = draft_cands[i] + tk_n[i] += 1 + if p == c[0]: + tk_cover1[i] += 1 + if p in c[:2]: + tk_cover2[i] += 1 + if p in c[:3]: + tk_cover3[i] += 1 + if profile: + t_verify += time.perf_counter() - _tic + _tic = time.perf_counter() + + accepted_len = min(matched + 1, step_k) + + if verify_snaps is not None: + # step 模式：逐 token 解码路径产生的快照可精确 direct commit。 + committed = commit_verified_snapshot(main_cache, verify_snaps, + accepted_len, verified_len=K) + else: + # batch 模式：一次并行验证后按接受长度直接提交，零 replay。 + # - 可裁剪 cache（KVCache）：trim 掉 rejected 后缀即精确。 + # - 递归状态 cache（Qwen3-Next gated-delta 的 conv/ssm）：verify 前向里用 + # multistate kernel 捕获的 per-token checkpoint 与 baseline kernel 逐 bit 等价 + # （见 mtp/kv_cache.py 与 tests/test_gated_delta_multistate.py），可精确直提交。 + # commit_verified_prefix 一次处理混合 cache：KV 走 trim、Arrays 走 checkpoint 直提交。 + # 若任一递归 cache 缺 checkpoint（异常情况）则返回 False，自动回退到下方 replay。 + committed = commit_verified_prefix(main_cache, verified_len=step_k, + accepted_len=accepted_len) + if snap_d is not None: + _restore(mtp_cache, snap_d) + + accepted_in = verify_in[:, :accepted_len] + if profile: + t_commit += time.perf_counter() - _tic + _tic = time.perf_counter() + if committed: + direct_commits += 1 + # 直接使用验证前向产生的 hidden 同步 MTP cache,避免主模型 replay。 + rH = vH[:, :accepted_len, :] + if matched == step_k: + new_tokens = drafts[:step_k] + x = drafts[-1] + else: + new_tokens = drafts[:matched] + [preds[matched]] + x = preds[matched] + else: + fallback_replays += 1 + if snap_m is None: + # 理论不可达:batch 直接提交路径 checkpoint 齐全,commit 必成功。若真走到这里, + # 说明前提被破坏(如未捕获 checkpoint),此时无快照可回退,直接抛错定位而非静默错算。 + raise RuntimeError("batch verify commit failed but snapshot was skipped") + # fallback:回滚到验证前,重放 accepted prefix,保证 ArraysCache 正确。 + _restore(main_cache, snap_m) + accepted_in = mx.array([[x0] + drafts[:matched]]) + replayed_tokens += accepted_in.shape[1] + rlogits, rH = forward_with_hidden(model, accepted_in, main_cache) + mx.eval(rlogits) + bonus = int(mx.argmax(rlogits[:, -1, :])) + new_tokens = drafts[:matched] + [bonus] + x = bonus + if profile: + t_replay += time.perf_counter() - _tic + _tic = time.perf_counter() + + _n_before = len(produced) + for t in new_tokens: + produced.append(t) + if len(produced) >= max_tokens: + break + # 只把「本步真正写入 produced」的 token 交给回调:多 token 步命中 max_tokens + # 上限时,截断掉的尾巴不应上报给流式消费者(避免超报被丢弃的 token)。 + _stop = on_tokens is not None and on_tokens(produced[_n_before:]) + + H_last = rH[:, -1:, :] + if mtp_cache is not None and hasattr(drafter, "sync"): + drafter.sync(prev_H_last, rH, accepted_in, mtp_cache) + if profile: + t_sync += time.perf_counter() - _tic + _tic = time.perf_counter() + n_steps += 1 + mx.eval(x, H_last) + if profile: + t_finalize += time.perf_counter() - _tic + if _stop: + break + + produced = produced[:max_tokens] + wall = time.perf_counter() - t0 + # cache 实际驻留 token 数(以可裁剪 cache 的 offset 为准),供跨轮复用精确对账: + # 末步多 token 跨 max_tokens 会 over-commit(cache 领先于 produced),据此识别并禁用复用。 + resident_tokens = None + for c in main_cache: + if getattr(c, "is_trimmable", None) and c.is_trimmable() and hasattr(c, "offset"): + resident_tokens = int(c.offset) + break + stats = { + "steps": n_steps, + "tokens": len(produced), + "resident_tokens": resident_tokens, + "avg_accept_len": round(len(produced) / max(n_steps, 1), 3), + "wall_s": round(wall, 3), + "verify_mode": verify_mode, + "direct_commits": direct_commits, + "fallback_replays": fallback_replays, + "replayed_tokens": replayed_tokens, + "accept_hist": accept_hist, # [恰好命中0,1,...,K 个草稿的步数] + "tree_rescues": tree_rescues, # 最小树第1 位成功救回步数(tree_top2 开时) + "tree_rescues_p1": tree_rescues_p1, # 最小树第2 位成功救回步数(tree_top2 开时) + } + if topk_probe > 0: + stats["topk_probe"] = { + "topk": topk_probe, "n": tk_n, + "cover_top1": tk_cover1, "cover_top2": tk_cover2, "cover_top3": tk_cover3, + } + if profile: + seg = t_draft + t_snap + t_verify + t_commit + t_replay + t_sync + t_finalize + stats.update({ + "t_draft_s": round(t_draft, 3), + "t_snap_s": round(t_snap, 3), + "t_verify_s": round(t_verify, 3), + "t_commit_s": round(t_commit, 3), + "t_replay_s": round(t_replay, 3), + "t_sync_s": round(t_sync, 3), + "t_finalize_s": round(t_finalize, 3), + # 「重放免费」上限:草稿+验证(不含快照,因 per-token 检查点设计无需全 cache 快照) + "proj_no_replay_tps": round(len(produced) / max(t_draft + t_verify, 1e-6), 2), + "measured_seg_tps": round(len(produced) / max(seg, 1e-6), 2), + }) + if ids_mode: + return produced, stats + return tok.decode(produced), stats diff --git a/mlx_streaming/mtp/kv_cache.py b/mlx_streaming/mtp/kv_cache.py new file mode 100644 index 0000000..f1c2aa4 --- /dev/null +++ b/mlx_streaming/mtp/kv_cache.py @@ -0,0 +1,276 @@ +"""MTP 投机解码的 KV/递归状态 cache 校验机制：per-token checkpoint + 快照/恢复/提交。 + +speculative decoding 需要在「验证 K 个草稿」后只保留 accepted_len 个 token 对 cache 的 +贡献。两类 cache 处理方式不同： +- 可裁剪 cache（KVCache）：直接 trim 掉 rejected 后缀。 +- 递归状态 cache（Qwen3-Next 线性注意力 ArraysCache）：必须在 verify 前向中逐 token 记录 + `[conv, ssm]` checkpoint 才能精确提交，否则回退到「快照 + replay accepted prefix」。 +""" +import mlx.core as mx +import mlx.nn as nn + +_QWEN3NEXT_CHECKPOINTS_PATCHED = False +_EMPTY_CACHE = object() + + +def enable_qwen3next_speculative_checkpoints(): + """给 Qwen3-Next 线性注意力层加 verify-time per-token cache checkpoint。 + + 普通前向仍走 mlx-lm 原实现;只有 `begin_speculative_checkpoints()` 标记过的 + ArraysCache 会走逐 token gated-delta ops,并把每个 prefix 后的 `[conv, ssm]` + 状态写入 `cache._spec_checkpoints`。 + """ + global _QWEN3NEXT_CHECKPOINTS_PATCHED + if _QWEN3NEXT_CHECKPOINTS_PATCHED: + return + + from mlx_lm.models.qwen3_next import Qwen3NextGatedDeltaNet + from mlx_streaming.core.linear_attn.gated_delta_multistate import ( + gated_delta_update_multistate, + ) + + orig_call = Qwen3NextGatedDeltaNet.__call__ + + def patched_call(self, inputs, mask=None, cache=None): + capture = bool(getattr(cache, "_capture_spec_checkpoints", False)) + if not capture: + return orig_call(self, inputs, mask=mask, cache=cache) + + B, S, _ = inputs.shape + q, k, v, z, b, a = self.fix_query_key_value_ordering( + self.in_proj_qkvz(inputs), self.in_proj_ba(inputs) + ) + + if cache is not None and cache[0] is not None: + conv_state = cache[0] + else: + conv_state = mx.zeros( + (B, self.conv_kernel_size - 1, self.conv_dim), + dtype=inputs.dtype, + ) + + mixed_qkv = mx.concatenate( + [q.reshape(B, S, -1), k.reshape(B, S, -1), v.reshape(B, S, -1)], axis=-1 + ) + if mask is not None: + mixed_qkv = mx.where(mask[..., None], mixed_qkv, 0) + conv_input = mx.concatenate([conv_state, mixed_qkv], axis=1) + + n_keep = self.conv_kernel_size - 1 + conv_checkpoints = [ + mx.contiguous(conv_input[:, i:i + n_keep, :]) + for i in range(1, S + 1) + ] + if cache is not None: + cache[0] = conv_checkpoints[-1] + + conv_out = nn.silu(self.conv1d(conv_input)) + q, k, v = [ + t.reshape(B, S, h, d) + for t, h, d in zip( + mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1), + [self.num_k_heads, self.num_k_heads, self.num_v_heads], + [self.head_k_dim, self.head_k_dim, self.head_v_dim], + ) + ] + + state = cache[1] if cache else None + 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) + + # 关键：用 multistate kernel 一次前向算出每个 token 处理后的 ssm 递归态。 + # 它与 baseline 解码走的mlx gated_delta_kernel 是同一份 kernel、同序、fp32， + # 因此 states_out[:, i] 与「逐 token 单步解码」逐 bit 等价（见 + # tests/test_gated_delta_multistate.py）——这是验证后能直接提交、零 replay 的根基。 + # 注意：kernel 路径内部处理 GQA（hk_idx 映射），不在 Python 侧 repeat q/k， + # 与 baseline kernel 路径对齐（旧 ops 路径的 repeat 会引入数值差异）。 + out, state, states_out = gated_delta_update_multistate( + q, k, v, a, b, self.A_log, self.dt_bias, state, mask + ) + + if cache is not None: + cache[1] = state + cache.advance(S) + cache._spec_checkpoints = [ + [conv_checkpoints[i], states_out[:, i]] + for i in range(S) + ] + cache._capture_spec_checkpoints = False + + out = self.norm(out, z) + return self.out_proj(out.reshape(B, S, -1)) + + Qwen3NextGatedDeltaNet.__call__ = patched_call + _QWEN3NEXT_CHECKPOINTS_PATCHED = True + + +def begin_speculative_checkpoints(caches): + """标记 ArraysCache 在下一次 verify forward 中记录 per-token checkpoint。""" + for c in caches: + if not c.is_trimmable(): + c._capture_spec_checkpoints = True + c._spec_checkpoints = None + + +# ----------------------------------------------------------------- cache 快照 +def _copy_state(st): + """深拷贝 cache.state(支持 None / array / list / tuple 嵌套)。""" + if st is None: + return None + if isinstance(st, (list, tuple)): + return type(st)(_copy_state(s) for s in st) + return mx.array(st) + + +def _iter_arrays(st): + if st is None: + return + if isinstance(st, (list, tuple)): + for s in st: + yield from _iter_arrays(s) + else: + yield st + + +def _snapshot(caches): + """对每个 cache 深拷贝 state + meta_state(强制 eval,防 update_and_fetch 原地改写)。""" + snaps = [] + arrays = [] + for c in caches: + if hasattr(c, "keys") and hasattr(c, "empty") and c.empty(): + snaps.append((_EMPTY_CACHE, c.meta_state)) + continue + st_copy = _copy_state(c.state) + arrays.extend(_iter_arrays(st_copy)) + snaps.append((st_copy, c.meta_state)) + if arrays: + mx.eval(arrays) + return snaps + + +def _restore(caches, snaps): + for c, (st, meta) in zip(caches, snaps): + if st is _EMPTY_CACHE: + if hasattr(c, "keys") and hasattr(c, "values"): + c.keys = None + c.values = None + c.offset = 0 + elif hasattr(c, "cache"): + c.cache = [None] * len(c.cache) + c.meta_state = meta + continue + # 必须装入副本:ArraysCache.state setter 是别名赋值(self.cache=v),后续前向 + # 的 cache[idx]=new 会原地改写这个 list 的元素,从而污染快照本体。救回路径会对 + # 同一 snap_m 恢复两次(救回时一次 + fallback 一次),若不复制,第二次恢复到的将是 + # 被中间前向污染的状态 → 递归 cache 损坏、输出发散。 + c.state = _copy_state(st) + c.meta_state = meta + + +def commit_verified_prefix(caches, verified_len: int, accepted_len: int) -> bool: + """把验证前向产生的 cache 直接提交到 accepted prefix。 + + vLLM 语义下,验证 K 个 token 后只保留 accepted_len 个 token 对 cache 的贡献。 + 对 KVCache 这类可裁剪 cache,直接 trim 掉 rejected 后缀即可;对 ArraysCache 这类 + 递归状态 cache,必须有 per-token checkpoint 才能精确提交,否则返回 False 让调用方 + 走 fallback replay。 + """ + rejected = verified_len - accepted_len + if rejected < 0: + raise ValueError("accepted_len cannot exceed verified_len") + def can_commit(c): + if c.is_trimmable(): + return True + checkpoints = getattr(c, "_spec_checkpoints", None) + return checkpoints is not None and len(checkpoints) >= accepted_len + + if not all(can_commit(c) for c in caches): + return False + for c in caches: + if c.is_trimmable(): + if rejected: + c.trim(rejected) + else: + c.state = c._spec_checkpoints[accepted_len - 1] + if hasattr(c, "_spec_checkpoints"): + c._spec_checkpoints = None + if hasattr(c, "_capture_spec_checkpoints"): + c._capture_spec_checkpoints = False + return True + + +# ----------------------------------------------------------- 树形验证 batch 工具 +def _tile_state(st, P: int): + """把 cache.state 沿 batch 轴(axis0)复制 P 份((1,...) → (P,...))。""" + if st is None: + return None + if isinstance(st, (list, tuple)): + return type(st)(_tile_state(s, P) for s in st) + return mx.contiguous(mx.repeat(st, P, axis=0)) + + +def _row_state(st, w: int): + """取 cache.state 的第 w 行((P,...) → (1,...)),保持 batch 维。""" + if st is None: + return None + if isinstance(st, (list, tuple)): + return type(st)(_row_state(s, w) for s in st) + return mx.contiguous(st[w:w + 1]) + + +def tile_caches(caches, P: int): + """把 batch=1 的 cache 全部平铺到 batch=P,供 batch-of-paths 树验证并行前向。 + + KVCache 的 state setter 会按数组 shape 复位 offset;ArraysCache 无 offset。两者 meta + (lengths/left_padding) 在贪婪解码里恒为 None,tile 后语义不变。前提:所有 cache 当前 batch=1。 + """ + for c in caches: + c.state = _tile_state(c.state, P) + + +def commit_tree_row(caches, verified_len: int, accepted_len: int, row: int) -> bool: + """把 batched(batch=P)验证前向的第 `row` 条路径,按 accepted_len 提交回 batch=1 主 cache。 + + 等价于「先按接受长度裁剪 rejected 后缀,再抽出赢家路径那一行」: + - 可裁剪 cache(KVCache):trim(rejected) 后取第 row 行 KV。 + - 递归状态 cache(ArraysCache):取 verify 前向捕获的 per-token checkpoint[accepted_len-1] + 的第 row 行 [conv, ssm]。要求 batched 前向前调过 begin_speculative_checkpoints。 + 提交后主 cache 变回 batch=1,后续解码逐 token 从赢家路径续。 + """ + rejected = verified_len - accepted_len + if rejected < 0: + raise ValueError("accepted_len cannot exceed verified_len") + + def can_commit(c): + if c.is_trimmable(): + return True + cks = getattr(c, "_spec_checkpoints", None) + return cks is not None and len(cks) >= accepted_len + + if not all(can_commit(c) for c in caches): + return False + for c in caches: + if c.is_trimmable(): + if rejected: + c.trim(rejected) + c.state = _row_state(c.state, row) + else: + ck = c._spec_checkpoints[accepted_len - 1] # [conv(P,...), ssm(P,...)] + c.state = [_row_state(ck[0], row), _row_state(ck[1], row)] + if hasattr(c, "_spec_checkpoints"): + c._spec_checkpoints = None + if hasattr(c, "_capture_spec_checkpoints"): + c._capture_spec_checkpoints = False + return True + + +def commit_verified_snapshot(caches, snapshots, accepted_len: int, + verified_len: int | None = None) -> bool: + """把 cache 恢复到 stepwise verify 的 accepted_len 后快照。""" + if verified_len is not None and accepted_len == verified_len: + return True + if accepted_len <= 0 or accepted_len > len(snapshots): + return False + _restore(caches, snapshots[accepted_len - 1]) + return True diff --git a/mlx_streaming/mtp/qwen3_next_mtp.py b/mlx_streaming/mtp/qwen3_next_mtp.py new file mode 100644 index 0000000..c47bd9a --- /dev/null +++ b/mlx_streaming/mtp/qwen3_next_mtp.py @@ -0,0 +1,95 @@ +"""Qwen3-Next MTP(多 token 预测)模块的 MLX 实现,复用 mlx-lm 现成子模块类。 + +前向(与 vLLM/sglang/trtllm 一致): + emb = pre_fc_norm_embedding(embed(next_id)) + hid = pre_fc_norm_hidden(主模型 last-layer hidden, norm 之前) + x = fc(concat([emb, hid], axis=-1)) # emb 在前 + x = layer(x) # 全注意力 + MoE 解码层(内部含残差) + logits = lm_head(norm(x)) +""" +from typing import Any, Callable, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_unflatten + +from mlx_lm.models.base import create_attention_mask +from mlx_lm.models.qwen3_next import ModelArgs, Qwen3NextDecoderLayer + + +class Qwen3NextMTP(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + h = args.hidden_size + eps = args.rms_norm_eps + self.embed_tokens = nn.Embedding(args.vocab_size, h) + self.pre_fc_norm_embedding = nn.RMSNorm(h, eps=eps) + self.pre_fc_norm_hidden = nn.RMSNorm(h, eps=eps) + self.fc = nn.Linear(2 * h, h, bias=False) + # layer_idx=3 配合 full_attention_interval=4 -> 全注意力 + MoE + self.layer = Qwen3NextDecoderLayer(args, layer_idx=3) + self.norm = nn.RMSNorm(h, eps=eps) + + def __call__( + self, + hidden: mx.array, # (B, L, H) 主模型 last-layer hidden(norm 前) + next_ids: mx.array, # (B, L) 每位置的"下一个" token id + lm_head: Callable[[mx.array], mx.array], + cache: Optional[Any] = None, + return_hidden: bool = False, + ): + emb = self.pre_fc_norm_embedding(self.embed_tokens(next_ids)) + hid = self.pre_fc_norm_hidden(hidden) + x = self.fc(mx.concatenate([emb, hid], axis=-1)) + mask = create_attention_mask(x, cache) if cache is not None else "causal" + x = self.layer(x, mask=mask, cache=cache) + H = self.norm(x) + logits = lm_head(H) + if return_hidden: + return logits, H + return logits + + +def mtp_step(mtp, hidden, token, lm_head, cache): + """单步:hidden(1,1,H) + token(1,1) -> (logits(1,V), mtp_hidden(1,1,H))。""" + logits, H = mtp(hidden, token, lm_head, cache=cache, return_hidden=True) + return logits[:, -1, :], H[:, -1:, :] + + +def mtp_advance(mtp, hidden, token, cache): + """只推进 MTP cache 并返回 hidden,不计算 lm_head logits。""" + emb = mtp.pre_fc_norm_embedding(mtp.embed_tokens(token)) + hid = mtp.pre_fc_norm_hidden(hidden) + x = mtp.fc(mx.concatenate([emb, hid], axis=-1)) + mask = create_attention_mask(x, cache) if cache is not None else "causal" + x = mtp.layer(x, mask=mask, cache=cache) + return mtp.norm(x) + + +def load_mtp(args: ModelArgs, weights_path: str, quantize: bool = True, + bits: int = 4) -> Qwen3NextMTP: + """加载抽取好的 MTP 权重(已 stack 专家、已 norm +1.0)到模块。 + + quantize=True 时按主模型约定量化:线性层用 `bits`-bit(默认 4),gate/shared_expert_gate + 恒用 8-bit。quantize=False 保持 bf16 全精度——草稿质量最高、接受率最高,但显存更大。 + 提高 bits(4→8)或 quantize=False 是直接抬升 MTP 草稿接受率的杠杆。 + """ + model = Qwen3NextMTP(args) + raw = mx.load(weights_path) + # 去掉 'mtp.' 前缀,映射到模块属性路径;mtp.layers.0.* -> layer.* + renamed = {} + for k, v in raw.items(): + nk = k[len("mtp."):] if k.startswith("mtp.") else k + nk = nk.replace("layers.0.", "layer.", 1) + renamed[nk] = v + model.update(tree_unflatten(list(renamed.items()))) + if quantize: + def pred(path, m): + if not hasattr(m, "to_quantized"): + return False + if path.endswith("mlp.gate") or path.endswith("shared_expert_gate"): + return {"group_size": 64, "bits": 8} + return True + nn.quantize(model, group_size=64, bits=bits, class_predicate=pred) + mx.eval(model.parameters()) + return model diff --git a/mlx_streaming/native_moe_ext.cpython-314-darwin.so b/mlx_streaming/native_moe_ext.cpython-314-darwin.so new file mode 100755 index 0000000..e9064e9 Binary files /dev/null and b/mlx_streaming/native_moe_ext.cpython-314-darwin.so differ diff --git a/mlx_streaming/prep/__init__.py b/mlx_streaming/prep/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mlx_streaming/prep/blob_layout.py b/mlx_streaming/prep/blob_layout.py new file mode 100644 index 0000000..78d8a1d --- /dev/null +++ b/mlx_streaming/prep/blob_layout.py @@ -0,0 +1,32 @@ +"""专家 blob 字节布局的单一真相源（v1 affine / v2 mxfp4）。 + +Segment 五元组 (proj, tensor, np_dtype_name, shape, nbytes)，与现有 +blob_loader._layout 的元组结构一致 → blob_loader 可直接消费、pack 取子集。 +""" +from typing import List, Tuple + +BLOB_V1_AFFINE = "expert_blob_v1" +BLOB_V2_MXFP4 = "expert_blob_v2_mxfp4" + +Segment = Tuple[str, str, str, Tuple[int, int], int] # proj, tensor, dtype, shape, nbytes + + +def _projs(hidden: int, inter: int) -> Tuple[Tuple[str, int, int], ...]: + return (("gate_proj", inter, hidden), ("up_proj", inter, hidden), ("down_proj", hidden, inter)) + + +def layout_for(fmt: str, hidden: int, inter: int, bits: int, group: int) -> Tuple[List[Segment], int]: + segs: List[Segment] = [] + for proj, out_d, in_d in _projs(hidden, inter): + words = in_d * bits // 32 + groups = in_d // group + segs.append((proj, "weight", "uint32", (out_d, words), out_d * words * 4)) + if fmt == BLOB_V1_AFFINE: + segs.append((proj, "scales", "uint16", (out_d, groups), out_d * groups * 2)) + segs.append((proj, "biases", "uint16", (out_d, groups), out_d * groups * 2)) + elif fmt == BLOB_V2_MXFP4: + segs.append((proj, "scales", "uint8", (out_d, groups), out_d * groups * 1)) + else: + raise ValueError(f"未知 blob format: {fmt}") + stride = sum(s[-1] for s in segs) + return segs, stride diff --git a/mlx_streaming/prep/download_retry.py b/mlx_streaming/prep/download_retry.py new file mode 100644 index 0000000..0b3d842 --- /dev/null +++ b/mlx_streaming/prep/download_retry.py @@ -0,0 +1,35 @@ +"""自动重试的断点续传下载：应对代理对大文件不稳定的情况。 + +反复调用 snapshot_download（huggingface_hub 会从 .incomplete 续传）， +每轮带超时；直到全部文件就绪或达到最大轮数。 +""" +import os +import sys +import time + +from huggingface_hub import snapshot_download + +REPO = os.environ.get("MODEL", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit") +MAX_TRIES = int(os.environ.get("MAX_TRIES", "200")) +os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "30") + + +def main(): + for i in range(1, MAX_TRIES + 1): + try: + p = snapshot_download( + REPO, + allow_patterns=["*.json", "*.txt", "*.safetensors", "tokenizer*", "*.model"], + max_workers=1, + ) + print(f"DONE -> {p}", flush=True) + return + except Exception as e: + print(f"[try {i}] 失败，续传重试: {repr(e)[:160]}", flush=True) + time.sleep(3) + print("EXHAUSTED: 达到最大重试轮数仍未完成", flush=True) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/extract_mtp.py b/mlx_streaming/prep/extract_mtp.py new file mode 100644 index 0000000..7454818 --- /dev/null +++ b/mlx_streaming/prep/extract_mtp.py @@ -0,0 +1,131 @@ +"""下载 Qwen3-Next 原版末分片并抽取/整理 MTP 权重为单文件。 + +MTP 全部权重在 model-00041-of-00041.safetensors(~3.9GB)。整理规则: + 1) 把 mlp.experts.{e}.{proj}.weight(512 个)stack 成 mlp.switch_mlp.{proj}.weight。 + 2) 对所有 RMSNorm 权重 +1.0(Qwen3-Next 用 zero-centered RMSNorm, + 与 mlx-lm sanitize 对主模型的处理一致;MTP 被 sanitize 过滤,故需自行补)。 +""" +import json +import os +import subprocess + +import mlx.core as mx + +# 与 mlx-lm qwen3_next.sanitize 一致的 norm 后缀(去掉 model.norm,因 MTP 用 mtp.norm) +_NORM_SUFFIXES = ( + ".input_layernorm.weight", + ".post_attention_layernorm.weight", + ".q_norm.weight", + ".k_norm.weight", +) + + +def _is_mtp_norm(key: str) -> bool: + if key.endswith(".pre_fc_norm_hidden.weight") or key.endswith( + ".pre_fc_norm_embedding.weight" + ): + return True + if key == "mtp.norm.weight": + return True + return any(key.endswith(sfx) for sfx in _NORM_SUFFIXES) + + +def bump_mtp_norms(weights: dict) -> dict: + out = {} + for k, v in weights.items(): + if _is_mtp_norm(k) and v.ndim == 1: + out[k] = v + 1.0 + else: + out[k] = v + return out + + +def stack_mtp_experts(weights: dict, num_experts: int) -> dict: + out = dict(weights) + prefix = "mtp.layers.0.mlp" + for proj in ("gate_proj", "up_proj", "down_proj"): + keys = [f"{prefix}.experts.{e}.{proj}.weight" for e in range(num_experts)] + if keys[0] not in out: + continue + stacked = mx.stack([out.pop(k) for k in keys]) + out[f"{prefix}.switch_mlp.{proj}.weight"] = stacked + return out + + +SHARD = "model-00041-of-00041.safetensors" +REPO = "Qwen/Qwen3-Next-80B-A3B-Instruct" +from mlx_streaming import config as _cfg +SHARD_DIR = os.environ.get("MTP_SHARD_DIR", "/tmp/qn_mtp_shard") # 临时分片，可留 /tmp +OUT_PATH = _cfg.mtp_out() # 默认 models/qn_mtp_weights.safetensors（持久） +CONFIG = _cfg.qn_config() # 默认 models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit/config.json + +_DL_URL = ( + f"https://www.modelscope.cn/api/v1/models/{REPO}/repo" + f"?Revision=master&FilePath={SHARD}" +) + + +def _expected_size() -> int: + """从 ModelScope 文件列表 API 取 SHARD 的真实字节数。 + + 不能用 HEAD/content-length:API 重定向只回小 JSON,content-length 不可靠。 + """ + url = ( + f"https://www.modelscope.cn/api/v1/models/{REPO}/repo/files" + f"?Revision=master&Recursive=true" + ) + out = subprocess.run( + ["curl", "-sL", "--max-time", "60", url], + capture_output=True, text=True, + ).stdout + data = json.loads(out) + files = data.get("Data", {}).get("Files") or data.get("Data", {}).get("files") or [] + for f in files: + name = f.get("Name") or f.get("Path") or f.get("name") + if name and SHARD in str(name): + return int(f.get("Size") or f.get("size") or 0) + return 0 + + +def download_shard() -> str: + os.makedirs(SHARD_DIR, exist_ok=True) + out = os.path.join(SHARD_DIR, SHARD) + expect = _expected_size() + for attempt in range(1, 9): + cur = os.path.getsize(out) if os.path.exists(out) else 0 + if expect and cur == expect: + print(f"SKIP shard already complete ({cur}B)") + return out + print(f"download attempt {attempt}: have {cur}B / {expect}B") + subprocess.run( + ["curl", "-sL", "--max-time", "3600", + "--speed-limit", "51200", "--speed-time", "30", + _DL_URL, "-o", out], + ) + cur = os.path.getsize(out) if os.path.exists(out) else 0 + if expect and cur != expect: + raise RuntimeError(f"download failed: {cur}B != {expect}B") + return out + + +def extract(shard_path: str, config_path: str, out_path: str) -> None: + with open(config_path) as f: + cfg = json.load(f) + num_experts = cfg["num_experts"] + all_w = mx.load(shard_path) + mtp = {k: v for k, v in all_w.items() if k.startswith("mtp.")} + print(f"原始 mtp 张量数: {len(mtp)}") + mtp = stack_mtp_experts(mtp, num_experts) + mtp = bump_mtp_norms(mtp) + mx.eval(list(mtp.values())) + mx.save_safetensors(out_path, mtp) + print(f"已写出 {len(mtp)} 个张量 -> {out_path}") + + +def main(): + shard = download_shard() + extract(shard, CONFIG, OUT_PATH) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/pack_blob_from_experts.py b/mlx_streaming/prep/pack_blob_from_experts.py new file mode 100644 index 0000000..1750830 --- /dev/null +++ b/mlx_streaming/prep/pack_blob_from_experts.py @@ -0,0 +1,87 @@ +"""直接把 per-expert safetensors 打包成「每专家一个连续 blob」(按层一个文件)。 + +字节布局由 prep/blob_layout.py 统一描述（v1 affine / v2 mxfp4），与 blob_loader 完全一致： +按 _split_meta.json 的 blob_format 选段表，每专家依次写各 proj 的 [weight, scales(, biases)]。 +- weight: uint32；v1 的 scales/biases 为 bf16 按 uint16 原始 2 字节；v2 mxfp4 的 scales 为 uint8。 + +环境变量：EXPERT_DIR(源 per-expert) / BLOB_DIR(输出) / BITS / GROUP / LAYERS(逗号或 all) +""" +import json +import os + +import mlx.core as mx +import numpy as np + +from mlx_streaming.prep.blob_layout import layout_for, BLOB_V1_AFFINE, BLOB_V2_MXFP4 + +EXPERT_DIR = os.environ.get("EXPERT_DIR", "/tmp/qwen3_next_experts_8bit_g128") +OUT = os.environ.get("BLOB_DIR", "/tmp/cb_8bit_blob") +BITS = int(os.environ.get("BITS", "8")) +GROUP = int(os.environ.get("GROUP", "128")) + + +def _meta(): + m = json.load(open(os.path.join(EXPERT_DIR, "_split_meta.json"))) + d = m["dims"] + fmt = m.get("blob_format", BLOB_V1_AFFINE) + return (int(d["num_experts"]), int(d["hidden"]), int(d["moe_intermediate"]), + fmt, d.get("quant_mode", "affine")) + + +def _raw_bytes(arr) -> bytes: + """uint32 → 4B/元素；uint8 → 1B/元素；其余(affine 的 bf16 scales/biases) → uint16 2B/元素。""" + if arr.dtype == mx.uint32: + return np.array(arr, copy=False).tobytes() + if arr.dtype == mx.uint8: + return np.array(arr, copy=False).tobytes() + return np.array(arr.view(mx.uint16), copy=False).tobytes() + + +def pack_layer(layer, num_experts, segs, stride, fmt, quant_mode) -> None: + out_path = os.path.join(OUT, f"layer{layer:02d}.blob") + with open(out_path, "wb") as f: + for e in range(num_experts): + w = mx.load(os.path.join(EXPERT_DIR, f"layer{layer:02d}_expert{e:03d}.safetensors")) + for proj, tensor, dt, shape, nb in segs: + b = _raw_bytes(w[f"{proj}.{tensor}"]) + assert len(b) == nb, f"L{layer} e{e} {proj}.{tensor}: {len(b)}!={nb}" + f.write(b) + assert os.path.getsize(out_path) == stride * num_experts + index = {"format": fmt, "quant_mode": quant_mode, "layer": layer, + "num_experts": num_experts, "stride": stride, + "page_aligned": stride % 16384 == 0, + "segments": [{"proj": p, "tensor": t, "nbytes": n} for p, t, _, _, n in segs]} + with open(os.path.join(OUT, f"layer{layer:02d}.blob.index.json"), "w") as f: + json.dump(index, f, ensure_ascii=False, indent=2) + + +def _resolve_layers(spec: str, num_experts: int) -> list: + spec = spec.strip() + if spec.lower() == "all": + layers = set() + for name in os.listdir(EXPERT_DIR): + if name.startswith("layer") and name.endswith("_expert000.safetensors"): + layers.add(int(name[5:7])) + return sorted(layers) + return [int(x) for x in spec.split(",") if x.strip()] + + +def main(): + os.makedirs(OUT, exist_ok=True) + num_experts, hidden, inter, fmt, quant_mode = _meta() + segs, stride = layout_for(fmt, hidden, inter, BITS, GROUP) + layers = _resolve_layers(os.environ.get("LAYERS", "all"), num_experts) + for i, L in enumerate(layers): + pack_layer(L, num_experts, segs, stride, fmt, quant_mode) + print(f" layer {L} packed ({i+1}/{len(layers)})", flush=True) + summary = {"format": fmt, "quant_mode": quant_mode, "stride": stride, + "page_aligned": stride % 16384 == 0, "num_experts": num_experts, + "layers": layers, "bits": BITS, "group_size": GROUP} + with open(os.path.join(OUT, "blob_index.json"), "w") as f: + json.dump(summary, f, ensure_ascii=False, indent=2) + print(json.dumps({"out": OUT, "n_layers": len(layers), "stride_bytes": stride, + "page_aligned": stride % 16384 == 0, "format": fmt}, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/pack_compute_buffers.py b/mlx_streaming/prep/pack_compute_buffers.py new file mode 100644 index 0000000..a8f2549 --- /dev/null +++ b/mlx_streaming/prep/pack_compute_buffers.py @@ -0,0 +1,83 @@ +"""把单层单 projection 打包成连续 compute buffers。 + +输出三个大 buffer: +- layer43.gate_proj.weight.bin +- layer43.gate_proj.scales.bin +- layer43.gate_proj.biases.bin + +这个格式用于验证“整层少数大 buffer + custom qlinear kernel 直接按 expert_id 读取”。 +""" +import argparse +import json +import os + +from mlx_streaming.prep.pack_expert_ranges import read_safetensors_payloads + + +def pack_projection(src_dir: str, out_dir: str, layer: int, proj: str, num_experts: int) -> dict: + os.makedirs(out_dir, exist_ok=True) + tensors = ["weight", "scales", "biases"] + handles = { + name: open(os.path.join(out_dir, f"layer{layer:02d}.{proj}.{name}.bin"), "wb") + for name in tensors + } + meta = { + "format": "mlx_streaming_compute_buffer_v1", + "src_dir": src_dir, + "layer": layer, + "proj": proj, + "num_experts": num_experts, + "tensors": {}, + } + try: + for expert in range(num_experts): + path = os.path.join(src_dir, f"layer{layer:02d}_expert{expert:03d}.safetensors") + payloads = {rec["key"]: rec for rec in read_safetensors_payloads(path)} + for name in tensors: + key = f"{proj}.{name}" + rec = payloads[key] + f = handles[name] + offset = f.tell() + f.write(rec["payload"]) + entry = meta["tensors"].setdefault(name, { + "dtype": rec["dtype"], + "shape_per_expert": rec["shape"], + "nbytes_per_expert": len(rec["payload"]), + "file": f"layer{layer:02d}.{proj}.{name}.bin", + "offsets": [], + }) + entry["offsets"].append(offset) + finally: + for f in handles.values(): + f.close() + meta_path = os.path.join(out_dir, f"layer{layer:02d}.{proj}.index.json") + with open(meta_path, "w") as f: + json.dump(meta, f, ensure_ascii=False, indent=2) + return meta + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--src", default=os.environ.get("EXPERT_DIR", "")) + ap.add_argument("--out", default=os.environ.get("COMPUTE_BUFFER_DIR", "")) + ap.add_argument("--layer", type=int, default=int(os.environ.get("LAYER", "43"))) + ap.add_argument("--proj", default=os.environ.get("PROJ", "gate_proj")) + args = ap.parse_args() + if not args.src: + raise SystemExit("--src / EXPERT_DIR required") + out_dir = args.out or os.path.join(args.src, "compute_buffers") + with open(os.path.join(args.src, "_split_meta.json")) as f: + model_meta = json.load(f) + num_experts = int(model_meta["dims"]["num_experts"]) + meta = pack_projection(args.src, out_dir, args.layer, args.proj, num_experts) + print(json.dumps({ + "out_dir": out_dir, + "layer": args.layer, + "proj": args.proj, + "num_experts": num_experts, + "files": {k: v["file"] for k, v in meta["tensors"].items()}, + }, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/pack_expert_bundles.py b/mlx_streaming/prep/pack_expert_bundles.py new file mode 100644 index 0000000..1d4edbd --- /dev/null +++ b/mlx_streaming/prep/pack_expert_bundles.py @@ -0,0 +1,70 @@ +"""把 per-expert safetensors 打包成 per-layer bundle。 + +输入目录: + layer00_expert000.safetensors + layer00_expert001.safetensors + ... + +输出目录: + layer_bundles/layer00.safetensors + +bundle 内 key: + expert000.gate_proj.weight + expert000.gate_proj.scales + ... + +运行时设置 EXPERT_BUNDLE=1 即可优先读取 bundle。 +""" +import json +import os +import time + +import mlx.core as mx + +SRC_DIR = os.environ.get("EXPERT_DIR", "/tmp/qwen3_next_experts") +OUT_DIR = os.environ.get("EXPERT_BUNDLE_DIR", os.path.join(SRC_DIR, "layer_bundles")) + + +def _pack_layer(src_dir: str, out_dir: str, layer: int, num_experts: int) -> str: + out = {} + for e in range(num_experts): + path = os.path.join(src_dir, f"layer{layer:02d}_expert{e:03d}.safetensors") + rec = mx.load(path) + prefix = f"expert{e:03d}." + for k, v in rec.items(): + out[prefix + k] = v + os.makedirs(out_dir, exist_ok=True) + out_path = os.path.join(out_dir, f"layer{layer:02d}.safetensors") + mx.save_safetensors(out_path, out) + return out_path + + +def main(): + t0 = time.perf_counter() + with open(os.path.join(SRC_DIR, "_split_meta.json")) as f: + meta = json.load(f) + layers = [int(x) for x in meta["moe_layers"]] + num_experts = int(meta["dims"]["num_experts"]) + os.makedirs(OUT_DIR, exist_ok=True) + for i, layer in enumerate(layers): + out_path = _pack_layer(SRC_DIR, OUT_DIR, layer, num_experts) + print(f" {i + 1}/{len(layers)} layer{layer:02d} -> {out_path} " + f"({round(time.perf_counter() - t0, 1)}s)", flush=True) + bundle_meta = { + "src_dir": SRC_DIR, + "out_dir": OUT_DIR, + "layers": layers, + "num_experts": num_experts, + "format": "expert{EEE}.{tensor_key}", + } + with open(os.path.join(OUT_DIR, "_bundle_meta.json"), "w") as f: + json.dump(bundle_meta, f, ensure_ascii=False, indent=2) + print(json.dumps({ + "out_dir": OUT_DIR, + "layers": len(layers), + "elapsed_s": round(time.perf_counter() - t0, 1), + }, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/pack_expert_ranges.py b/mlx_streaming/prep/pack_expert_ranges.py new file mode 100644 index 0000000..03567b7 --- /dev/null +++ b/mlx_streaming/prep/pack_expert_ranges.py @@ -0,0 +1,135 @@ +"""把 per-expert safetensors 的原始 payload 打包成 per-layer range pack。 + +这个格式只用于 de-risk：避免运行时解析大量小 safetensors header，先验证 +“按 expert byte range 读取”是否比多个 `mx.load` 更快。 +""" +import argparse +import json +import os +import struct +import time +from pathlib import Path + +ALIGN = 64 + + +def _align(n: int, align: int = ALIGN) -> int: + rem = n % align + return n if rem == 0 else n + (align - rem) + + +def read_safetensors_payloads(path: str) -> list[dict]: + """读取 safetensors header 和每个 tensor 的原始 payload bytes。""" + with open(path, "rb") as f: + header_len = struct.unpack(" dict: + """打包单层所有专家，返回 index dict。""" + os.makedirs(out_dir, exist_ok=True) + pack_path = os.path.join(out_dir, f"layer{layer:02d}.pack") + index_path = os.path.join(out_dir, f"layer{layer:02d}.index.json") + tsv_path = os.path.join(out_dir, f"layer{layer:02d}.idx") + tensors = [] + experts = [] + with open(pack_path, "wb") as pack: + for expert in range(num_experts): + src = os.path.join(src_dir, f"layer{layer:02d}_expert{expert:03d}.safetensors") + expert_start = _align(pack.tell()) + if expert_start > pack.tell(): + pack.write(b"\0" * (expert_start - pack.tell())) + for rec in read_safetensors_payloads(src): + offset = _align(pack.tell()) + if offset > pack.tell(): + pack.write(b"\0" * (offset - pack.tell())) + payload = rec["payload"] + pack.write(payload) + tensors.append({ + "layer": layer, + "expert_id": expert, + "key": rec["key"], + "dtype": rec["dtype"], + "shape": rec["shape"], + "offset": offset, + "nbytes": len(payload), + }) + expert_end = pack.tell() + experts.append({ + "layer": layer, + "expert_id": expert, + "offset": expert_start, + "nbytes": expert_end - expert_start, + }) + index = { + "format": "mlx_streaming_expert_pack_v1", + "alignment": ALIGN, + "layer": layer, + "num_experts": num_experts, + "pack": os.path.basename(pack_path), + "experts": experts, + "tensors": tensors, + } + with open(index_path, "w") as f: + json.dump(index, f, ensure_ascii=False, indent=2) + with open(tsv_path, "w") as f: + f.write("kind\tlayer\texpert_id\tkey\tdtype\tshape\toffset\tnbytes\n") + for rec in experts: + f.write( + f"EXPERT\t{rec['layer']}\t{rec['expert_id']}\t*\t*\t*\t" + f"{rec['offset']}\t{rec['nbytes']}\n" + ) + for rec in tensors: + shape = ",".join(str(x) for x in rec["shape"]) + f.write( + f"TENSOR\t{rec['layer']}\t{rec['expert_id']}\t{rec['key']}\t" + f"{rec['dtype']}\t{shape}\t{rec['offset']}\t{rec['nbytes']}\n" + ) + return index + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--src", default=os.environ.get("EXPERT_DIR", "/tmp/qwen3_next_experts")) + ap.add_argument("--out", default=os.environ.get("EXPERT_PACK_DIR")) + ap.add_argument("--layers", default=os.environ.get("LAYERS", "")) + args = ap.parse_args() + src = args.src + out = args.out or os.path.join(src, "layer_packs") + with open(os.path.join(src, "_split_meta.json")) as f: + meta = json.load(f) + all_layers = [int(x) for x in meta["moe_layers"]] + layers = [int(x) for x in args.layers.split(",") if x.strip()] if args.layers else all_layers + num_experts = int(meta["dims"]["num_experts"]) + t0 = time.perf_counter() + for i, layer in enumerate(layers): + pack_layer(src, out, layer, num_experts) + print(f" {i + 1}/{len(layers)} layer{layer:02d} ({round(time.perf_counter() - t0, 1)}s)", flush=True) + with open(os.path.join(out, "_pack_meta.json"), "w") as f: + json.dump({ + "src": src, + "out": out, + "layers": layers, + "num_experts": num_experts, + "format": "mlx_streaming_expert_pack_v1", + }, f, ensure_ascii=False, indent=2) + print(json.dumps({"out": out, "layers": len(layers)}, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/repack_expert_blobs.py b/mlx_streaming/prep/repack_expert_blobs.py new file mode 100644 index 0000000..5c5413a --- /dev/null +++ b/mlx_streaming/prep/repack_expert_blobs.py @@ -0,0 +1,92 @@ +"""把 compute buffer 重打包成"每专家一个连续 blob"格式（按层一个文件）。 + +blob 内顺序：gate[w,s,b] + up[w,s,b] + down[w,s,b]，正好 864KB（= 54×16KB 页对齐）。 +读一个专家 = 1 次 pread(stride, e*stride)，而不是当前的 9 次散读。 + +环境变量：COMPUTE_BUFFER_DIR(源)、BLOB_DIR(输出)、LAYERS(逗号分隔，默认 15,25) +""" +import json +import os + +import numpy as np + +HIDDEN = 2048 +INTER = 512 +GROUP = 128 +BITS = 2 +NUM_EXPERTS = 512 +SRC = os.environ.get("COMPUTE_BUFFER_DIR", "/tmp/cb_2bit_g128") +OUT = os.environ.get("BLOB_DIR", "/tmp/cb_2bit_blob") +PROJS = (("gate_proj", INTER, HIDDEN), ("up_proj", INTER, HIDDEN), ("down_proj", HIDDEN, INTER)) + + +def _layout(): + """返回 [(proj, tensor, nbytes_per_expert), ...] 和单专家总字节。""" + segs = [] + for proj, out_dim, in_dim in PROJS: + words = in_dim * BITS // 32 + groups = in_dim // GROUP + segs.append((proj, "weight", out_dim * words * 4)) + segs.append((proj, "scales", out_dim * groups * 2)) + segs.append((proj, "biases", out_dim * groups * 2)) + return segs, sum(s[2] for s in segs) + + +def repack_layer(layer: int) -> dict: + os.makedirs(OUT, exist_ok=True) + segs, stride = _layout() + # 源 .bin 以 uint8 原始字节映射，按专家字节偏移切片 + raw = {} + for proj, _, _ in PROJS: + base = os.path.join(SRC, f"layer{layer:02d}.{proj}") + for tensor in ("weight", "scales", "biases"): + raw[(proj, tensor)] = np.memmap(f"{base}.{tensor}.bin", dtype=np.uint8, mode="r") + out_path = os.path.join(OUT, f"layer{layer:02d}.blob") + with open(out_path, "wb") as f: + for e in range(NUM_EXPERTS): + for proj, tensor, nb in segs: + buf = raw[(proj, tensor)] + f.write(buf[e * nb:(e + 1) * nb].tobytes()) + assert os.path.getsize(out_path) == stride * NUM_EXPERTS + index = { + "format": "expert_blob_v1", + "layer": layer, + "num_experts": NUM_EXPERTS, + "stride": stride, + "page_aligned": stride % 16384 == 0, + "segments": [{"proj": p, "tensor": t, "nbytes": n} for p, t, n in segs], + } + with open(os.path.join(OUT, f"layer{layer:02d}.blob.index.json"), "w") as f: + json.dump(index, f, ensure_ascii=False, indent=2) + return index + + +def _resolve_layers(spec: str) -> list[int]: + """LAYERS=all → 读源 compute buffer 现有的全部层；否则按逗号解析。""" + spec = spec.strip() + if spec.lower() == "all": + layers = [] + for name in os.listdir(SRC): + if name.endswith(".gate_proj.weight.bin") and name.startswith("layer"): + layers.append(int(name[len("layer"):len("layer") + 2])) + return sorted(set(layers)) + return [int(x) for x in spec.split(",") if x.strip()] + + +def main(): + layers = _resolve_layers(os.environ.get("LAYERS", "15,25")) + _, stride = _layout() + for L in layers: + repack_layer(L) + # 汇总 index，供 loader 校验/发现 + summary = {"format": "expert_blob_v1", "stride": stride, + "page_aligned": stride % 16384 == 0, "num_experts": NUM_EXPERTS, + "layers": layers} + with open(os.path.join(OUT, "blob_index.json"), "w") as f: + json.dump(summary, f, ensure_ascii=False, indent=2) + print(json.dumps({"out": OUT, "n_layers": len(layers), "stride_bytes": stride, + "page_aligned": stride % 16384 == 0}, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/prep/split_experts.py b/mlx_streaming/prep/split_experts.py new file mode 100644 index 0000000..6326de3 --- /dev/null +++ b/mlx_streaming/prep/split_experts.py @@ -0,0 +1,69 @@ +"""路线 B 离线工具：把模型里堆叠的 switch_mlp 专家权重按专家拆成 per-expert 小文件。 + +拆分后每个文件 layer{L:02d}_expert{E:03d}.safetensors 含扁平 dict： + gate_proj.weight / gate_proj.scales / gate_proj.biases + up_proj.weight / up_proj.scales / up_proj.biases + down_proj.weight / down_proj.scales / down_proj.biases +（非量化模型则只有 .weight，可能还有 .bias） + +一次拆分、逐专家物化，全程低内存。 +""" +import os +import sys +import json + +import mlx.core as mx + +PROJ_NAMES = ["gate_proj", "up_proj", "down_proj"] + + +def split_switch_glu(switch_glu, out_dir: str, layer: int) -> int: + """把一个 SwitchGLU 的三组 SwitchLinear 沿专家维拆成 per-expert 文件。返回专家数。""" + os.makedirs(out_dir, exist_ok=True) + E = switch_glu.gate_proj.num_experts + for e in range(E): + d = {} + for proj_name in PROJ_NAMES: + proj = getattr(switch_glu, proj_name) + for pname, p in proj.parameters().items(): + if isinstance(p, mx.array) and p.ndim >= 1 and p.shape[0] == E: + d[f"{proj_name}.{pname}"] = p[e] + mx.eval(d) # 只物化这一个专家 + path = os.path.join(out_dir, f"layer{layer:02d}_expert{e:03d}.safetensors") + mx.save_safetensors(path, d) + return E + + +def split_model(model_path: str, out_dir: str) -> dict: + """加载模型（lazy）并把所有 MoE 层的专家拆到 out_dir。返回 {dims/统计}。""" + from mlx_lm import load + model, _ = load(model_path, lazy=True) + os.makedirs(out_dir, exist_ok=True) + moe_layers = [] + dims = None + for l, layer in enumerate(model.layers): + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): + sm = mlp.switch_mlp + split_switch_glu(sm, out_dir, l) + moe_layers.append(l) + if dims is None: + gp = sm.gate_proj + dims = { + "hidden": gp.input_dims, + "moe_intermediate": gp.output_dims, + "num_experts": gp.num_experts, + "group_size": getattr(gp, "group_size", None), + "bits": getattr(gp, "bits", None), + } + meta = {"out_dir": out_dir, "moe_layers": moe_layers, "dims": dims} + with open(os.path.join(out_dir, "_split_meta.json"), "w") as f: + json.dump(meta, f, ensure_ascii=False, indent=2) + return meta + + +if __name__ == "__main__": + mp = sys.argv[1] + od = sys.argv[2] if len(sys.argv) > 2 else "/tmp/mlx_qwen3_experts" + m = split_model(mp, od) + print(json.dumps(m, ensure_ascii=False, indent=2)) diff --git a/mlx_streaming/runtime/__init__.py b/mlx_streaming/runtime/__init__.py new file mode 100644 index 0000000..eaf1dad --- /dev/null +++ b/mlx_streaming/runtime/__init__.py @@ -0,0 +1 @@ +"""生产推理入口：基线/流式/投机解码等 run_* 启动脚本。""" diff --git a/mlx_streaming/runtime/run_baseline.py b/mlx_streaming/runtime/run_baseline.py new file mode 100644 index 0000000..17deeb5 --- /dev/null +++ b/mlx_streaming/runtime/run_baseline.py @@ -0,0 +1,42 @@ +"""全量驻留基线：正常加载 + 生成，记录 RSS/峰值/decode 速度。""" +import os +import time +import json + +import mlx.core as mx +from mlx_lm import load, generate + +from mlx_streaming.core.mem import snapshot, reset_peak + +MODEL = os.environ.get("MODEL", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit") +PROMPT = os.environ.get("PROMPT", "用三句话解释什么是混合专家模型。") +MAXTOK = int(os.environ.get("MAXTOK", "128")) + + +def main(): + reset_peak() + t0 = time.perf_counter() + model, tok = load(MODEL) # 默认全量加载 + mx.eval(model.parameters()) # 强制全部驻留，作为对照上界 + load_done = snapshot() + t1 = time.perf_counter() + + text = generate(model, tok, prompt=PROMPT, max_tokens=MAXTOK, verbose=False) + t2 = time.perf_counter() + after = snapshot() + + out = { + "mode": "baseline_resident", + "model": MODEL, + "load_s": round(t1 - t0, 2), + "gen_s": round(t2 - t1, 2), + "tok_per_s": round(MAXTOK / (t2 - t1), 2), + "rss_gb_after_load": round(load_done.rss_bytes / 1e9, 2), + "rss_gb_after_gen": round(after.rss_bytes / 1e9, 2), + "mlx_peak_gb": round(after.mlx_peak_bytes / 1e9, 2), + } + print(json.dumps(out, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/runtime/run_mmap.py b/mlx_streaming/runtime/run_mmap.py new file mode 100644 index 0000000..77ee85a --- /dev/null +++ b/mlx_streaming/runtime/run_mmap.py @@ -0,0 +1,69 @@ +"""路线 A：lazy（+ 若支持则 mmap）加载，可选设常驻/缓存上限，量内存与速度。 + +环境探针结论（mlx-lm 0.31.3）：load 支持 lazy，但不支持 use_mmap。 +本脚本用 _filtered_load_kwargs 只传该版本真正支持的参数，缺失的自动剔除。 +""" +import os +import time +import json +import inspect + +import mlx.core as mx +from mlx_lm import load, generate + +from mlx_streaming.core.mem import snapshot, reset_peak, clear_cache + +MODEL = os.environ.get("MODEL", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit") +PROMPT = os.environ.get("PROMPT", "用三句话解释什么是混合专家模型。") +MAXTOK = int(os.environ.get("MAXTOK", "128")) +WIRED_GB = os.environ.get("WIRED_GB") # 例如 "6"；空=不设 +CACHE_GB = os.environ.get("CACHE_GB") # 例如 "1"；空=不设 + + +def _maybe(fn_name, val_bytes): + fn = getattr(mx, fn_name, None) or getattr(getattr(mx, "metal", object()), fn_name, None) + if fn and val_bytes is not None: + prev = fn(int(val_bytes)) + print(f"{fn_name}({int(val_bytes)}) 旧值={prev}") + + +def _filtered_load_kwargs(): + sig = inspect.signature(load) + want = {"lazy": True, "use_mmap": True} + return {k: v for k, v in want.items() if k in sig.parameters} + + +def main(): + if WIRED_GB: + _maybe("set_wired_limit", float(WIRED_GB) * 1e9) + if CACHE_GB: + _maybe("set_cache_limit", float(CACHE_GB) * 1e9) + + reset_peak() + kw = _filtered_load_kwargs() + print("load kwargs:", kw) + t0 = time.perf_counter() + model, tok = load(MODEL, **kw) # 不强制 eval 全部参数！ + t1 = time.perf_counter() + after_load = snapshot() # 关键：此时 RSS 应远低于基线（专家还没换入） + + text = generate(model, tok, prompt=PROMPT, max_tokens=MAXTOK, verbose=False) + t2 = time.perf_counter() + clear_cache() + after_gen = snapshot() + + out = { + "mode": "mmap_lazy", "model": MODEL, + "wired_gb": WIRED_GB, "cache_gb": CACHE_GB, + "load_s": round(t1 - t0, 2), "gen_s": round(t2 - t1, 2), + "tok_per_s": round(MAXTOK / (t2 - t1), 2), + "rss_gb_after_load": round(after_load.rss_bytes / 1e9, 2), + "rss_gb_after_gen": round(after_gen.rss_bytes / 1e9, 2), + "mlx_active_gb_after_gen": round(after_gen.mlx_active_bytes / 1e9, 2), + "mlx_peak_gb": round(after_gen.mlx_peak_bytes / 1e9, 2), + } + print(json.dumps(out, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/runtime/run_mtp_spec.py b/mlx_streaming/runtime/run_mtp_spec.py new file mode 100644 index 0000000..746d850 --- /dev/null +++ b/mlx_streaming/runtime/run_mtp_spec.py @@ -0,0 +1,409 @@ +"""真实 80B MTP 自投机基准 + 与非投机贪婪逐 token 一致性校验。 + +环境变量:K / MAXTOK / PROMPT / QN_CONFIG / MTP_OUT(其余主模型路径见 validate_mtp)。 +""" +import json +import os +import statistics +import time + +import mlx.core as mx +from mlx_lm.models.qwen3_next import ModelArgs + +from mlx_streaming.core.moe import native_moe +from mlx_streaming.core.mem import snapshot, reset_peak +from mlx_streaming.mtp.drafter import MTPDrafter +from mlx_streaming.mtp.generate import forward_with_hidden, mtp_generate, prefill_chunked +from mlx_streaming.mtp.qwen3_next_mtp import load_mtp +from mlx_streaming.model_builder import build_streaming_model + +from mlx_streaming import config as _cfg +QN_CONFIG = _cfg.qn_config() +MTP_OUT = _cfg.mtp_out() +PROMPT = os.environ.get("PROMPT", "用三句话解释什么是混合专家模型。") +MAXTOK = int(os.environ.get("MAXTOK", "96")) +K = int(os.environ.get("K", "3")) +PIN_HOT = int(os.environ.get("PIN_HOT", "0")) +PIN_CAL_TOK = int(os.environ.get("PIN_CAL_TOK", "32")) +# 稳态测速:warmup 跑满 MAXTOK(热 baseline+spec 两条路径的 Metal kernel 与常驻专家池/预取), +# 再各重复 REPEAT 次取中位数,避免冷启动编译/补池污染绝对 tok/s。WARMUP_TOK=0 关 warmup。 +WARMUP_TOK = int(os.environ.get("WARMUP_TOK", str(MAXTOK))) +REPEAT = int(os.environ.get("REPEAT", "3")) + + +def _spec_once(model, drafter, tok, n, k): + ids, stats = mtp_generate(model, drafter, tok, + mx.array([tok.encode(PROMPT)]), + n, K=k, ids_mode=True, profile=True) + tps = round(stats["tokens"] / stats["wall_s"], 2) + return ids, stats, tps + + +def _baseline_greedy(model, tok, prompt, n): + cache = model.make_cache() + ids = mx.array([tok.encode(prompt)]) + t0 = time.perf_counter() + # prefill 分块:把整段 prompt 的激活峰值压到与 decode 同稳态(见 config.prefill_chunk)。 + logits, _ = prefill_chunked(model, ids, cache) + _dump_margin = bool(os.environ.get("DUMP_MARGIN")) + out = [] + for _ in range(n): + lg = logits[:, -1, :] + nxt = int(mx.argmax(lg)) + out.append(nxt) + if _dump_margin: + # 诊断:每步 top-2 logit 与差值,用于判定路径间发散是「FP 近平局」还是「真错」。 + v = lg.reshape(-1) + top2 = mx.argpartition(-v, 2)[:2] + t2 = [int(x) for x in top2.tolist()] + vv = {t: float(v[t]) for t in t2} + st = sorted(t2, key=lambda t: -vv[t]) + print(f"MARGIN step={len(out)-1} top1={st[0]}({vv[st[0]]:.5f}) " + f"top2={st[1]}({vv[st[1]]:.5f}) gap={vv[st[0]]-vv[st[1]]:.6f}", flush=True) + cur = mx.array([[nxt]]) + mx.eval(cur) + if len(out) >= n: + break + logits, _ = forward_with_hidden(model, cur, cache) + return out, round(n / (time.perf_counter() - t0), 2) + + +def main(): + reset_peak() + model, tok, store = build_streaming_model() + with open(QN_CONFIG) as f: + args = ModelArgs.from_dict(json.load(f)) + mtp = load_mtp(args, MTP_OUT, quantize=True) + mtp.embed_tokens = model.model.embed_tokens # 共享主模型 embedding + drafter = MTPDrafter(mtp, model.lm_head) + + # warmup:同时跑 baseline + spec 两条路径,编译 Metal kernel(含 multistate/batch verify) + # 并把专家常驻池/预取热起来,确保后续测的是稳态而非冷启动。默认 warmup=MAXTOK 跑满全长。 + if WARMUP_TOK > 0: + _baseline_greedy(model, tok, PROMPT, WARMUP_TOK) + _spec_once(model, drafter, tok, WARMUP_TOK, K) + + # ---- baseline 稳态:重复 REPEAT 次取中位数;最后一次清零统计供命中率口径 ---- + base = None + base_tps_runs = [] + for r in range(REPEAT): + if r == REPEAT - 1: + store.reset_stats() + base, bt = _baseline_greedy(model, tok, PROMPT, MAXTOK) + base_tps_runs.append(bt) + base_tps = statistics.median(base_tps_runs) + base_miss, base_hit = store.misses, store.hits + base_prefetch_loads = store._resident.prefetch_loads + base_prefetch_hits = store._resident.prefetch_hits + + # 可选:baseline 之后、spec 之前校准每层热专家，并预取/钉入 resident pool。 + # 这样 disk_load_ratio 的分母仍是未 pin baseline，能直接衡量 pin 是否压低 spec miss。 + if PIN_HOT > 0: + store.record = True + _baseline_greedy(model, tok, PROMPT, PIN_CAL_TOK) + for li in store.recorded_layers(): + store.pin(li, store.hot(li, PIN_HOT)) + store.record = False + store.reset_stats() + + # 诊断:开 handler 触发时刻探针(仅覆盖 spec 阶段,enable 会清零旧日志)。 + _hprof = bool(os.environ.get("STAGING_HPROF")) + if _hprof: + from mlx_streaming import native_moe_ext as _Nhp + _Nhp.staging_hprof_enable(True) + + # ---- spec 稳态:重复 REPEAT 次取中位数;最后一次清零统计 + reset_peak + 保留其 stats ---- + ids, stats, spec_tps_runs = None, None, [] + for r in range(REPEAT): + if r == REPEAT - 1: + store.reset_stats() + reset_peak() + from mlx_streaming.core.profiling import tprof_reset, union_reset + tprof_reset() # 探针只统计最终测量轮(与命中率/内存口径一致) + union_reset() # 并集专家数也只统计最终轮 + ids, stats, st = _spec_once(model, drafter, tok, MAXTOK, K) + spec_tps_runs.append(st) + spec_tps = statistics.median(spec_tps_runs) + spec_miss, spec_hit = store.misses, store.hits + spec_prefetch_loads = store._resident.prefetch_loads + spec_prefetch_hits = store._resident.prefetch_hits + after = snapshot() + proj = stats.get("proj_no_replay_tps", 0.0) + result = { + "K": K, + "max_tokens": MAXTOK, + "warmup_tok": WARMUP_TOK, + "repeat": REPEAT, + "exact_match": ids == base, + "n_mismatch": sum(1 for a, b in zip(ids, base) if a != b), + "avg_accept_len": stats["avg_accept_len"], + "steps": stats["steps"], + "verify_mode": stats.get("verify_mode"), + "direct_commits": stats.get("direct_commits"), + "fallback_replays": stats.get("fallback_replays"), + "replayed_tokens": stats.get("replayed_tokens"), + "spec_tok_per_s": spec_tps, + "baseline_tok_per_s": base_tps, + "speedup": round(spec_tps / max(base_tps, 1e-6), 2), + "spec_tps_runs": spec_tps_runs, + "baseline_tps_runs": base_tps_runs, + "spec_tps_minmax": [min(spec_tps_runs), max(spec_tps_runs)], + # 分段计时与「重放免费」投机上限 + "t_draft_s": stats.get("t_draft_s"), + "t_snap_s": stats.get("t_snap_s"), + "t_verify_s": stats.get("t_verify_s"), + "t_commit_s": stats.get("t_commit_s"), + "t_replay_s": stats.get("t_replay_s"), + "t_sync_s": stats.get("t_sync_s"), + "t_finalize_s": stats.get("t_finalize_s"), + "proj_no_replay_tps": proj, + "proj_no_replay_speedup": round(proj / max(base_tps, 1e-6), 2), + "baseline_disk_loads": base_miss, + "baseline_prefetch_loads": base_prefetch_loads, + "baseline_prefetch_hits": base_prefetch_hits, + "baseline_hit_rate": round(base_hit / max(base_hit + base_miss, 1), 3), + "spec_disk_loads": spec_miss, + "spec_prefetch_loads": spec_prefetch_loads, + "spec_prefetch_hits": spec_prefetch_hits, + "spec_hit_rate": round(spec_hit / max(spec_hit + spec_miss, 1), 3), + "disk_load_ratio": round(spec_miss / max(base_miss, 1), 2), + # 双源 acquire 分路计数:n_miss==0 走全 GPU 快路径;任一路由 miss 则整层落 host 慢路径 + # (.tolist 全批同步 + demand 读盘)。fallback 占比高 → 即使 hit 高,慢路径仍按"层"频繁触发。 + "gpu_fastpath": getattr(store._resident, "gpu_fastpath", None), + "gpu_fallback": getattr(store._resident, "gpu_fallback", None), + "pin_hot": PIN_HOT, + "pin_cal_tok": PIN_CAL_TOK, + "pinned_experts": store.pinned_count(), + "expert_slots": store.capacity, + "mlx_active_gb": round(after.mlx_active_bytes / 1e9, 2), + "mlx_peak_gb": round(after.mlx_peak_bytes / 1e9, 2), + "rss_gb": round(after.rss_bytes / 1e9, 2), + # 内存分块(清缓冲后的真实常驻):A 权重 / B 专家池 / C staging / D MTP / F 激活 + "mem_breakdown": _mem_breakdown(model, store, mtp), + "prefill_chunk": _cfg.prefill_chunk(), + "native_stage_cache": native_moe.stage_cache_stats(), + "bg_stats": (store._bg.stats() if getattr(store, "_bg", None) is not None else None), + "window_prof": _window_prof(), + "predict_recall": _predict_recall(), + "miss_attrib": _miss_attrib(), + "prefetch_tprof": _prefetch_tprof(stats.get("wall_s")), + "union_experts": _union_prof(), + } + # 噪声地板测量口径:DUMP_IDS=1 时把 baseline greedy 与 spec 的完整 token 序列打进日志, + # 供跨进程 run-to-run 逐位对比(默认关闭,不污染常规输出)。 + if os.environ.get("DUMP_IDS"): + print("DUMP_BASE_IDS " + json.dumps(list(base))) + print("DUMP_SPEC_IDS " + json.dumps(list(ids))) + # 字节真值校验自证口径:开 STG_VERIFY 时,把两处校验器的累计计数 + # (ok/bad/calls)打进日志。关键:让「0 BAD」可判真伪——若 calls==0 说明本配置根本 + # 没触发该校验器(如 STG_VERIFY 在 zerocopy_dual 路径不接线),此时 0 BAD 是空结论。 + if os.environ.get("STG_VERIFY"): + from mlx_streaming.core.cache import resident_pool as _rp_mod + from mlx_streaming.core.cache import virtual_pool as _vp_mod + _vsum = { + "STG_VERIFY.resident(verify_acquire_bytes)": dict(_rp_mod._stg_verify_state), + "STG_VERIFY.virtual(_verify_native_bytes)": dict(_vp_mod._stg_verify_state), + } + print("VERIFY_SUMMARY " + json.dumps(_vsum, ensure_ascii=False)) + print(json.dumps(result, ensure_ascii=False, indent=2)) + if _hprof: + # dump (gen, layer, t_fire) 原始日志,供离线分析回调触发时刻分布。 + _flat = _Nhp.staging_hprof_get() # 扁平 [gen,layer,t, ...] + _path = os.environ.get("STAGING_HPROF_OUT", "/tmp/ab/hprof.jsonl") + os.makedirs(os.path.dirname(_path), exist_ok=True) + _n = 0 + with open(_path, "w") as _f: + for i in range(0, len(_flat), 3): + _f.write(json.dumps([int(_flat[i]), int(_flat[i + 1]), float(_flat[i + 2])]) + "\n") + _n += 1 + print(f"[hprof] wrote {_n} handler records to {_path}") + + +def _tree_nbytes(obj): + from mlx.utils import tree_flatten + return sum(v.nbytes for _, v in tree_flatten(obj) if isinstance(v, mx.array)) + + +def _mem_breakdown(model, store, mtp): + """把 decode 稳态内存拆成各块(清掉 MLX 可回收缓冲后量真实常驻)。 + + A 常驻非专家权重 / B 专家常驻池 / C staging 预取 / D MTP drafter(扣共享 embedding)/ + F 激活+临时(= active底 − A−B−C−D)。返回 GiB 字典。 + """ + GIB = 1024 ** 3 + A = _tree_nbytes(model.parameters()) + rp = store._resident + B = sum(v.nbytes for pool in rp._pools.values() for v in pool.values()) + C = 0 + stg = getattr(store, "_staging", None) + if stg is not None: + from mlx.utils import tree_flatten + C = sum(v.nbytes for _, v in tree_flatten(getattr(stg, "__dict__", {})) + if isinstance(v, mx.array)) + # D:MTP drafter 全部权重减去与主模型共享的 embedding(run 里 mtp.embed_tokens = 主模型的) + D = _tree_nbytes(mtp.parameters()) + emb = getattr(getattr(mtp, "embed_tokens", None), "weight", None) + shared = 0 + if isinstance(emb, mx.array): + # 共享 embedding 在 A 已计入;若 mtp 的 embed 与主模型同一对象则从 D 扣除避免重复 + shared = _tree_nbytes(mtp.embed_tokens) if hasattr(mtp, "embed_tokens") else 0 + D = max(0, D - shared) + mx.clear_cache() + active = mx.get_active_memory() + peak = mx.get_peak_memory() + F = max(0, active - (A + B + C + D)) + pool_rows = sum(rp.allocated_slots(l) for l in rp._pools) + return { + "A_weights_gib": round(A / GIB, 3), + "B_expert_pool_gib": round(B / GIB, 3), + "B_pool_rows": pool_rows, + "C_staging_gib": round(C / GIB, 3), + "D_mtp_drafter_gib": round(D / GIB, 3), + "F_activation_temp_gib": round(F / GIB, 3), + "active_live_gib": round(active / GIB, 3), + "peak_gib": round(peak / GIB, 3), + } + + +def _union_prof(): + """按 seq 分桶汇报每层路由专家并集大小(avg/max/p99/分层);seq=K 桶即 MTP verify 的专家并集。 + + 需 UNION_PROF=1 才有数据。返回 {seq: {avg_union, max_union, p99_union, per_layer...}}; + verify 桶 = seq==K(verify_in=[x, d_1..d_{K-1}] 恰 K 个 token,见 mtp/generate.py)。 + U_max 即该桶各层并集的全局最大值,是池 cap 下限的正确性依据(cap 必须 ≥ U_max 才保证不溢出)。 + """ + from mlx_streaming.core.profiling import UNION_PROF as U, UNION_SAMPLES as S + if not U: + return None + + def _p(vals, q): + # 最近秩法分位数(vals 已升序):q∈[0,1]。样本少时直接给保守上界。 + if not vals: + return None + idx = max(0, min(len(vals) - 1, int(round(q * (len(vals) - 1))))) + return vals[idx] + + out = {} + for seq in sorted(U): + s, n = U[seq] + rec = {"avg_union": round(s / max(1, n), 2), "n_layer_calls": n} + # 由分层原始样本汇总出该 seq 桶的 U_max / p99（cap 下限的正确性依据）。 + layer_map = S.get(seq) + if layer_map: + flat = sorted(v for lst in layer_map.values() for v in lst) + rec["max_union"] = flat[-1] + rec["p99_union"] = _p(flat, 0.99) + rec["p50_union"] = _p(flat, 0.50) + rec["n_samples"] = len(flat) + # 分层分布：每层的 max / avg / 样本数（按层号排序）。 + per_layer = {} + for li in sorted(layer_map): + lst = layer_map[li] + per_layer[li] = { + "max": max(lst), + "avg": round(sum(lst) / len(lst), 2), + "n": len(lst), + } + rec["per_layer"] = per_layer + rec["max_layer_idx"] = max(per_layer, key=lambda li: per_layer[li]["max"]) + out[f"seq{seq}"] = rec + # 便捷:verify 桶 = seq==K(verify_in=[x, d_1..d_{K-1}] 恰 K 个 token) + verify_seq = K + verify = out.get(f"seq{verify_seq}") + return {"by_seq": out, "verify_seq": verify_seq, + "verify_avg_union": (verify or {}).get("avg_union"), + "verify_max_union": (verify or {}).get("max_union"), + "verify_p99_union": (verify or {}).get("p99_union")} + + +def _miss_attrib(): + from mlx_streaming.core.profiling import MISS_ATTRIB as M + if not M["n"]: + return None + routed = max(1, M["routed"]) + miss = M["miss_A_predicted"] + M["miss_B_unpredicted"] + out = { + "routed": M["routed"], + "hit_rate": round(M["resident_hit"] / routed, 4), + "miss_rate": round(miss / routed, 4), + # A：预测到却没进池（budget/时序/驱逐）；B：没预测到（召回缺口） + "miss_A_predicted": M["miss_A_predicted"], + "miss_B_unpredicted": M["miss_B_unpredicted"], + "A_share_of_miss": round(M["miss_A_predicted"] / max(1, miss), 4), + "B_share_of_miss": round(M["miss_B_unpredicted"] / max(1, miss), 4), + } + if M["dec_n"]: + # decode/verify 热路径专用（与上面 prefill-主导的全量分开看）： + dr = max(1, M["dec_routed"]) + dmiss = M["dec_miss_A"] + M["dec_miss_B"] + out["decode"] = { + "routed": M["dec_routed"], + "hit_rate": round(M["dec_resident_hit"] / dr, 4), + "miss_rate": round(dmiss / dr, 4), + "miss_A_predicted": M["dec_miss_A"], + "miss_B_unpredicted": M["dec_miss_B"], + "A_share_of_miss": round(M["dec_miss_A"] / max(1, dmiss), 4), + "B_share_of_miss": round(M["dec_miss_B"] / max(1, dmiss), 4), + # miss_A 细分:时序(pread 没完成) vs 驱逐(就绪过但 acquire 前不在池) + "miss_A_timing": M["dec_miss_A_timing"], + "miss_A_evicted": M["dec_miss_A_evicted"], + "A_timing_share": round(M["dec_miss_A_timing"] / max(1, M["dec_miss_A"]), 4), + "A_evicted_share": round(M["dec_miss_A_evicted"] / max(1, M["dec_miss_A"]), 4), + } + return out + + +def _predict_recall(): + from mlx_streaming.core.profiling import PREDICT_RECALL_PROF as P + if not P["n"]: + return None + return {"recall": round(P["hit"] / max(1, P["routed"]), 4), + "avg_routed": round(P["routed"] / P["n"], 2), "n": P["n"]} + + +def _window_prof(): + from mlx_streaming.core.profiling import WINDOW_PROF + if not WINDOW_PROF["n"]: + return None + return {"avg_ms": round(WINDOW_PROF["sum_s"] / WINDOW_PROF["n"] * 1000, 4), + "n": WINDOW_PROF["n"]} + + +def _prefetch_tprof(wall_s=None): + """预取 host 墙钟探针汇总(PREFETCH_TPROF=1 才有数据)。 + + 汇报每段:总秒数、占最终测量轮 wall 的百分比、单次调用均值(ms)。 + 注意:这是"主线程不可重叠的 host 时间";gate matmul / pool scatter 的 GPU 执行落在 + 前向末尾统一 eval,不在此口径,需用消融(NATIVE_NO_SUBMIT/NATIVE_NO_PROMOTE)差值测。 + """ + from mlx_streaming.core.profiling import PREFETCH_TPROF as T + if not (T["predict_n"] or T["promote_n"] or T["submit_n"]): + return None + + def _seg(s, n): + d = {"total_s": round(T[s], 4)} + if wall_s: + d["pct_wall"] = round(T[s] / wall_s * 100, 2) + if n: + d["avg_ms"] = round(T[s] / n * 1000, 4) + return d + + host_total = T["predict_s"] + T["submit_s"] + T["promote_s"] + out = { + "predict": _seg("predict_s", T["predict_n"]), + "submit": _seg("submit_s", T["submit_n"]), + "promote": _seg("promote_s", T["promote_n"]), + # promote 内部细分(take 锁读 / route 成员同步 / place 切片+scatter 入图) + "promote_take": _seg("take_s", T["promote_n"]), + "promote_route": _seg("route_s", T["promote_n"]), + "promote_place": _seg("place_s", T["promote_n"]), + "host_total_s": round(host_total, 4), + "host_total_pct_wall": (round(host_total / wall_s * 100, 2) if wall_s else None), + "counts": {"predict_n": T["predict_n"], "submit_n": T["submit_n"], + "promote_n": T["promote_n"], "place_experts": T["place_experts"]}, + } + return out + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/runtime/run_spec.py b/mlx_streaming/runtime/run_spec.py new file mode 100644 index 0000000..01abf79 --- /dev/null +++ b/mlx_streaming/runtime/run_spec.py @@ -0,0 +1,116 @@ +"""投机解码 + 流式专家 验证：target=流式 Qwen3-30B，draft=常驻 Qwen3-0.6B。 + +测「独立 draft 投机」在 I/O 流式 MoE 上到底有没有加速：对比不投机 vs 不同 +num_draft_tokens 的 decode tok/s 与草稿接受率。 + +环境变量： + MODEL/EXPERT_DIR/EXPERT_BITS/EXPERT_GROUP/EXPERT_SLOTS 同 run_streaming + DRAFT draft 模型路径（默认 /tmp/qwen3_draft） + NDRAFTS 要扫的 num_draft_tokens 列表（逗号分隔，默认 "2,3,4"） + MAXTOK/PROMPT +""" +import os +import time +import json + +import mlx.core as mx +from mlx_lm import load, stream_generate + +from mlx_streaming.core.mem import snapshot, reset_peak +from mlx_streaming.core.cache.expert_store import FileExpertStore +from mlx_streaming.core.prefetch.patch import patch_model_filebacked + +MODEL = os.environ.get("MODEL", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit") +EXPERT_DIR = os.environ.get("EXPERT_DIR", "/tmp/mlx_qwen3_experts_2bit") +EXPERT_BITS = int(os.environ.get("EXPERT_BITS", "2")) +EXPERT_GROUP = int(os.environ.get("EXPERT_GROUP", "64")) +EXPERT_SLOTS = int(os.environ.get("EXPERT_SLOTS", "64")) +DRAFT = os.environ.get("DRAFT", "/tmp/qwen3_draft") +NDRAFTS = [int(x) for x in os.environ.get("NDRAFTS", "2,3,4").split(",")] +MAXTOK = int(os.environ.get("MAXTOK", "128")) +PROMPT = os.environ.get("PROMPT", "用三句话解释什么是混合专家模型。") + + +def _first_moe_dims(model): + for layer in model.layers: + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): + gp = mlp.switch_mlp.gate_proj + return {"hidden": gp.input_dims, "moe_inter": gp.output_dims} + raise RuntimeError("无 MoE 层") + + +def _build_target(): + model, tok = load(MODEL, lazy=True) + dims = _first_moe_dims(model) + # 混合精度：从专家目录 meta 读 proj_bits/bits/group（有则优先于环境变量默认） + bits, group, proj_bits, layer_proj_bits = EXPERT_BITS, EXPERT_GROUP, None, None + meta_path = os.path.join(EXPERT_DIR, "_split_meta.json") + if os.path.exists(meta_path): + with open(meta_path) as f: + ed = json.load(f).get("dims", {}) + bits = ed.get("bits", bits) + group = ed.get("group_size", group) + proj_bits = ed.get("proj_bits") + if "per_layer_proj_bits" in ed: + layer_proj_bits = {int(k): v for k, v in ed["per_layer_proj_bits"].items()} + store = FileExpertStore(EXPERT_DIR, capacity=EXPERT_SLOTS) + patch_model_filebacked(model, store, dims["hidden"], dims["moe_inter"], + group, bits, proj_bits=proj_bits, + layer_proj_bits=layer_proj_bits) + return model, tok, store, {"bits": bits, "group": group, "proj_bits": proj_bits, + "layered": layer_proj_bits is not None} + + +def _run(model, tok, draft_model, num_draft, store): + store.reset_stats() if hasattr(store, "reset_stats") else None + n_tok = 0 + n_draft_acc = 0 + last_tps = 0.0 + kwargs = {} + if draft_model is not None: + kwargs["num_draft_tokens"] = num_draft + t0 = time.perf_counter() + for r in stream_generate(model, tok, prompt=PROMPT, max_tokens=MAXTOK, + draft_model=draft_model, **kwargs): + n_tok = r.generation_tokens + last_tps = r.generation_tps + if getattr(r, "from_draft", False): + n_draft_acc += 1 + wall = time.perf_counter() - t0 + return { + "num_draft_tokens": num_draft if draft_model is not None else None, + "gen_tokens": n_tok, + "tok_per_s": round(last_tps, 2), + "wall_s": round(wall, 2), + "draft_accepted_frac": round(n_draft_acc / n_tok, 3) if n_tok else 0.0, + "expert_hit_rate": round(store.hit_rate(), 3), + } + + +def main(): + reset_peak() + model, tok, store, qinfo = _build_target() + draft_model, _ = load(DRAFT) # 常驻小 draft + mx.eval(draft_model.parameters()) + + results = [] + # 1) 不投机基线 + results.append({"mode": "no_spec", **_run(model, tok, None, 0, store)}) + # 2) 投机：扫 num_draft_tokens + for nd in NDRAFTS: + results.append({"mode": "spec", **_run(model, tok, draft_model, nd, store)}) + + after = snapshot() + print(json.dumps({ + "target": MODEL, "draft": DRAFT, + "expert_bits": qinfo["bits"], "expert_proj_bits": qinfo["proj_bits"], + "expert_slots": EXPERT_SLOTS, + "rss_gb": round(after.rss_bytes / 1e9, 2), + "mlx_peak_gb": round(after.mlx_peak_bytes / 1e9, 2), + "results": results, + }, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/runtime/run_streaming.py b/mlx_streaming/runtime/run_streaming.py new file mode 100644 index 0000000..7e55d5b --- /dev/null +++ b/mlx_streaming/runtime/run_streaming.py @@ -0,0 +1,180 @@ +"""路线 B 端到端：lazy 加载 + 离线拆分专家 + 文件后端流式 + generate，量内存/速度/命中率。 + +环境变量： + MODEL 模型 repo 或本地路径（默认 models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit） + EXPERT_DIR per-expert 拆分输出目录（默认 /tmp/mlx_qwen3_experts） + EXPERT_SLOTS 每层 LRU 专家槽数（默认 8=top_k），worst-case 常驻≈槽数×MoE层数 + PROMPT/MAXTOK 提示与生成 token 数 + WIRED_GB 可选，set_wired_limit 上限（GB） + CACHE_GB 可选，set_cache_limit 上限（GB），控制 MLX 缓冲复用上界 +""" +import os +import time +import json +import statistics + +import mlx.core as mx +from mlx_lm import load, generate + +from mlx_streaming import config +from mlx_streaming.core.mem import snapshot, reset_peak, clear_cache +from mlx_streaming.core.cache.expert_store import FileExpertStore +from mlx_streaming.prep.split_experts import split_model +from mlx_streaming.core.prefetch.patch import patch_model_filebacked +from mlx_streaming.model_builder import load_pool_profile + +MODEL = os.environ.get("MODEL", "models/Qwen3-Next-80B-A3B-Instruct-MLX-8bit") +EXPERT_DIR = os.environ.get("EXPERT_DIR", "/tmp/mlx_qwen3_experts") +EXPERT_SLOTS = int(os.environ.get("EXPERT_SLOTS", "8")) # 每层槽数 +PROMPT = os.environ.get("PROMPT", "用三句话解释什么是混合专家模型。") +MAXTOK = int(os.environ.get("MAXTOK", "128")) +# 稳态测速:warmup 跑满 MAXTOK(把整段会路由到的专家都装进常驻池 + 编译 Metal kernel), +# 再重复 REPEAT 次取中位数,避免冷启动/补池污染绝对 tok/s。WARMUP_TOK=0 关 warmup。 +WARMUP_TOK = int(os.environ.get("WARMUP_TOK", str(MAXTOK))) +REPEAT = int(os.environ.get("REPEAT", "3")) +WIRED_GB = os.environ.get("WIRED_GB") +CACHE_GB = os.environ.get("CACHE_GB") +CLEAR_ON_EVICT = os.environ.get("CLEAR_ON_EVICT", "0") == "1" +PIN_HOT = int(os.environ.get("PIN_HOT", "0")) # ③ 每层钉住的热专家数（0=关） +CAL_TOK = int(os.environ.get("CAL_TOK", "32")) # 校准生成 token 数 +# 专家量化格式覆盖：当 EXPERT_DIR 指向重量化目录(2/3-bit)时，文件后端 QSL 要用对应 bit/group +EXPERT_BITS = os.environ.get("EXPERT_BITS") +EXPERT_GROUP = os.environ.get("EXPERT_GROUP") + +def _first_moe_dims(model): + for layer in model.layers: + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): + gp = mlp.switch_mlp.gate_proj + return { + "hidden": gp.input_dims, "moe_inter": gp.output_dims, + "group_size": getattr(gp, "group_size", 64), + "bits": getattr(gp, "bits", 4), + } + raise RuntimeError("模型里没有找到 MoE 层") + + +def main(): + if WIRED_GB: + fn = getattr(mx, "set_wired_limit", None) + if fn: + print("set_wired_limit 旧值=", fn(int(float(WIRED_GB) * 1e9))) + if CACHE_GB: + fn = getattr(mx, "set_cache_limit", None) + if fn: + print("set_cache_limit 旧值=", fn(int(float(CACHE_GB) * 1e9))) + + # 1. 拆分专家到磁盘（若尚未拆分） + if not os.path.exists(os.path.join(EXPERT_DIR, "_split_meta.json")): + print("拆分专家到", EXPERT_DIR, "...") + t = time.perf_counter() + meta = split_model(MODEL, EXPERT_DIR) + print("拆分完成", round(time.perf_counter() - t, 1), "s; MoE 层=", len(meta["moe_layers"])) + + reset_peak() + # 2. lazy 加载（不强制 eval 全部） + t0 = time.perf_counter() + model, tok = load(MODEL, lazy=True) + dims = _first_moe_dims(model) + # 专家目录 meta 是 bit/group/proj_bits 的权威来源（重量化目录都会写）；先用它覆盖， + # 再让 EXPERT_BITS/EXPERT_GROUP 环境变量做最高优先级手动覆盖。proj_bits 非空走混合精度。 + proj_bits = None + layer_proj_bits = None + meta_path = os.path.join(EXPERT_DIR, "_split_meta.json") + if os.path.exists(meta_path): + with open(meta_path) as f: + ed = json.load(f).get("dims", {}) + dims["bits"] = ed.get("bits", dims["bits"]) + dims["group_size"] = ed.get("group_size", dims["group_size"]) + proj_bits = ed.get("proj_bits") + if "per_layer_proj_bits" in ed: # 逐层混合：键转回 int 层号 + layer_proj_bits = {int(k): v for k, v in ed["per_layer_proj_bits"].items()} + if EXPERT_BITS: + dims["bits"] = int(EXPERT_BITS) + if EXPERT_GROUP: + dims["group_size"] = int(EXPERT_GROUP) + # 每层池预算 profile:默认从 {EXPERT_DIR}/pool_profile.json 自动启用(无损省内存)， + # EXPERT_POOL_PROFILE 可显式指定路径或设 none 关闭。命中率/输出/吞吐不变。 + layer_caps = load_pool_profile(EXPERT_DIR) + # 3. 文件后端 patch，丢弃常驻堆叠 switch_mlp + store = FileExpertStore(EXPERT_DIR, capacity=EXPERT_SLOTS, layer_caps=layer_caps, + clear_on_evict=CLEAR_ON_EVICT, record=PIN_HOT > 0) + n = patch_model_filebacked(model, store, dims["hidden"], dims["moe_inter"], + dims["group_size"], dims["bits"], + proj_bits=proj_bits, layer_proj_bits=layer_proj_bits) + t1 = time.perf_counter() + after_patch = snapshot() + + # 分块 prefill:把 mlx_lm.generate 内部 prefill 步长压到 config.prefill_chunk()(默认 2), + # 整段 prefill 的激活峰值 ∝prompt 长度 → ∝chunk,使 prefill 与 decode 同稳态。 + # PREFILL_CHUNK=0 时不传,回退 mlx_lm 默认 2048。 + _ps = config.prefill_chunk() + _gen_kw = {"prefill_step_size": _ps} if _ps > 0 else {} + + # 3.5 ③ 热专家常驻：先校准跑一遍统计激活频率，钉住每层最热的 PIN_HOT 个专家 + cal_s = 0.0 + if PIN_HOT > 0: + tc = time.perf_counter() + generate(model, tok, prompt=PROMPT, max_tokens=CAL_TOK, verbose=False, **_gen_kw) + for li in store.recorded_layers(): + store.pin(li, store.hot(li, PIN_HOT)) + store.record = False + store.reset_stats() + cal_s = round(time.perf_counter() - tc, 2) + + # 3.6 warmup:跑满 MAXTOK 把整段专家装进常驻池并编译 Metal kernel(PIN_HOT 已校准时 + # 这里仍补满全长,确保第一次正式测量即稳态)。 + if WARMUP_TOK > 0: + generate(model, tok, prompt=PROMPT, max_tokens=WARMUP_TOK, verbose=False, **_gen_kw) + + # 4. 生成:重复 REPEAT 次取中位数;最后一次清零专家统计供命中率口径对应稳态。 + reset_peak() + text = None + gen_runs = [] + for r in range(REPEAT): + if r == REPEAT - 1: + store.reset_stats() + tg = time.perf_counter() + text = generate(model, tok, prompt=PROMPT, max_tokens=MAXTOK, verbose=False, **_gen_kw) + gen_runs.append(round(MAXTOK / (time.perf_counter() - tg), 2)) + tok_per_s = statistics.median(gen_runs) + t2 = time.perf_counter() + clear_cache() + after_gen = snapshot() + + out = { + "mode": "streaming_filebacked", "model": MODEL, + "expert_dir": EXPERT_DIR, + "expert_bits": dims["bits"], "expert_group": dims["group_size"], + "proj_bits": proj_bits, "layered": layer_proj_bits is not None, + "per_layer_slots": EXPERT_SLOTS, "pool_profile": bool(layer_caps), + "wired_gb": WIRED_GB, "cache_gb": CACHE_GB, + "clear_on_evict": CLEAR_ON_EVICT, + "pin_hot": PIN_HOT, "cal_s": cal_s, + "warmup_tok": WARMUP_TOK, "repeat": REPEAT, + "patched_moe_layers": n, + "resident_experts": store.resident_count(), + "pinned_experts": store.pinned_count(), + "load_patch_s": round(t1 - t0, 2), + "tok_per_s": tok_per_s, + "tok_per_s_runs": gen_runs, + "tok_per_s_minmax": [min(gen_runs), max(gen_runs)], + "rss_gb_after_patch": round(after_patch.rss_bytes / 1e9, 2), + "rss_gb_after_gen": round(after_gen.rss_bytes / 1e9, 2), + "mlx_active_gb_after_gen": round(after_gen.mlx_active_bytes / 1e9, 2), + "mlx_peak_gb": round(after_gen.mlx_peak_bytes / 1e9, 2), + "expert_hit_rate": round(store.hit_rate(), 4), + "expert_hits": store.hits, "expert_misses": store.misses, + # GPU remap 取证:整层全命中走 GPU 快路径 vs 有 miss 回退 host 的层调用占比 + "gpu_fastpath": store._resident.gpu_fastpath, + "gpu_fallback": store._resident.gpu_fallback, + "gpu_fastpath_frac": round( + store._resident.gpu_fastpath + / max(1, store._resident.gpu_fastpath + store._resident.gpu_fallback), 4), + "sample": text[:240], + } + print(json.dumps(out, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/server/__init__.py b/mlx_streaming/server/__init__.py new file mode 100644 index 0000000..65b0769 --- /dev/null +++ b/mlx_streaming/server/__init__.py @@ -0,0 +1,9 @@ +"""OpenAI 兼容的 FastAPI server。 + +引擎复用 TUI 的 ChatBackend 抽象:SPARKLE_FAKE=1(或测试显式注入)时用 FakeBackend +免模型开发/测试;默认用 MLXBackend 装配真实流式 MoE 引擎。 +入口:python -m mlx_streaming.server +""" +from mlx_streaming.server.app import create_app + +__all__ = ["create_app"] diff --git a/mlx_streaming/server/__main__.py b/mlx_streaming/server/__main__.py new file mode 100644 index 0000000..2d03b9e --- /dev/null +++ b/mlx_streaming/server/__main__.py @@ -0,0 +1,95 @@ +"""sparkle OpenAI 兼容 server 入口:python -m mlx_streaming.server + +模型路径参数与 sparkle chat 一致(默认值同样来自 mlx_streaming/config.py 的环境变量)。 +SPARKLE_FAKE=1 时用 FakeBackend,免模型开发/测试。 +""" +import argparse +import os + +from mlx_streaming import config + + +def _build_parser(): + p = argparse.ArgumentParser( + prog="sparkle-server", + description="sparkle:OpenAI 兼容的流式 MoE 推理 server") + p.add_argument("--host", default="127.0.0.1", help="监听地址(默认 127.0.0.1)") + p.add_argument("--port", type=int, default=8317, help="监听端口(默认 8317)") + p.add_argument("--model", default=config.model_path(), help="主模型路径(MLX 量化)") + p.add_argument("--expert-dir", default=config.expert_dir(), + help="拆分后的 per-expert 目录") + p.add_argument("--mtp-out", default=config.mtp_out(), help="MTP 权重文件") + p.add_argument("--qn-config", default=config.qn_config(), + help="Qwen3-Next 配置 JSON") + p.add_argument("-k", "--k", type=int, default=3, help="MTP 投机宽度(默认 3)") + p.add_argument("-n", "--max-tokens", type=int, default=4096, + help="每轮最多生成的新 token 数(默认 4096)") + p.add_argument("--expert-slots", type=int, default=32, + help="常驻专家池容量(默认 32,同时作为侧区行数默认)") + p.add_argument("--spec-slots", type=int, default=None, + help="侧区行数 POOL_SPEC_SLOTS(默认跟随 --expert-slots)") + return p + + +def _safe_prefill_chunk(args) -> int: + """按专家池实际可安放容量算 prefill 步长上限，取 min(启动方给的值, 上限)。 + + 一次前向路由的唯一专家最多 chunk × top_k 个，真实区能安放的只有 + expert_slots − AUTOPIN 钉死数（钉死的永不驱逐）。超出会被 block.py 分流到 + host/fetch 路径：结果仍正确，但每层每 chunk 重读盘，prefill 慢几倍。 + 读不到 top_k 时按 Qwen3-Next 的 10 兜底（偏保守，宁可 chunk 小）。 + """ + import json + + top_k = 10 + try: + with open(args.qn_config, encoding="utf-8") as f: + top_k = max(1, int(json.load(f).get("num_experts_per_tok", top_k))) + except Exception: # noqa: BLE001 配置读不到/字段缺失都退到保守默认 + pass + slots = max(1, int(args.expert_slots)) + pinned = int(slots * config.autopin_budget_frac()) if config.autopin() else 0 + limit = max(1, (slots - pinned) // top_k) + want = int(os.environ.get("PREFILL_CHUNK") or limit) + if want > limit: + print(f"[PREFILL_CHUNK] {want} → {limit}：{slots} 槽扣掉 AUTOPIN 钉死 {pinned} 后只能安放 " + f"{slots - pinned} 个专家，chunk×top_k({top_k}) 不得超过它，否则 prefill 退到读盘慢路径。", + flush=True) + return min(want, limit) + + +def main(argv=None): + args = _build_parser().parse_args(argv) + + # KV_QUANT 维持 Sparkle 原始配方(开):此前"不按规则回答"已证实是 tools 未渲染所致, + # 与 KV 量化无关(且其通过质量门:token 一致率≥95%、logits cosine≥0.99)。 + # KV_QUANT=0 会使 prefill 慢 ~3 倍(fp16 KV 带宽),业务长提示下代价不可接受。 + # 显式 export KV_QUANT=0 可关闭做质量对照。 + + # 前缀快照头部的**上限**(实际长度由引擎按 system+tools 的真实边界自动取,见 + # MLXBackend._fixed_prefix_len)。写死数字曾是主要的慢因:业务 system+4 工具实测 + # 3972 token,固定切 2048 只覆盖一半,每个新会话都要白 prefill 1900+ token + # (prefill 约 30 tok/s,即多等一分钟)。这里只兜住极端情况下快照过大。 + os.environ.setdefault("PREFIX_SNAPSHOT_HEAD", "8192") + + # AUTOPIN 用一半槽位钉历史热专家(冷启动不慢热),另一半留给 demand 动态换入。 + os.environ.setdefault("AUTOPIN_BUDGET_FRAC", "0.5") + + # PREFILL_CHUNK 必须由槽位推导,不能让启动方拍数字:每个 chunk 的一次前向最多路由 + # chunk × top_k 个唯一专家,而真实区能安放的只有 expert_slots − AUTOPIN 钉死数。超了 + # 会被分流到 host/fetch 路径(正确但每层重读盘,慢几倍),槽位调小后写死的 chunk 就失效。 + os.environ["PREFILL_CHUNK"] = str(_safe_prefill_chunk(args)) + + import uvicorn + + from mlx_streaming.server.app import create_app + + # timeout_graceful_shutdown:SIGTERM 后最多等 3 秒就强制关闭连接。 + # 外部客户端的轮询持有 keep-alive 长连接,不设超时 uvicorn 会永远 + # 卡在 "Waiting for connections to close",进程不退、内存不释放。 + uvicorn.run(create_app(args), host=args.host, port=args.port, + timeout_graceful_shutdown=3) + + +if __name__ == "__main__": + main() diff --git a/mlx_streaming/server/admin.py b/mlx_streaming/server/admin.py new file mode 100644 index 0000000..52ecb8d --- /dev/null +++ b/mlx_streaming/server/admin.py @@ -0,0 +1,116 @@ +"""极简 admin 页:单文件 HTML + 内联 JS,不引框架/CDN。中文文案,2.5s 轮询 stats。""" + +ADMIN_HTML = """ + + + + +Sparkle 引擎控制台 + + + +

Sparkle 引擎控制台

+ +
+
引擎状态-
+
模型-
+
最近一轮吞吐- tok/s
+
进程内存 RSS-
+
+ +
+
+ 常驻专家池 expert_slots + + 32 +
+
预估池内存:(slots × ~160MB(g64));修改后引擎后台重建,期间对话不可用
+
+ 投机宽度 k + +
+
+ 最大生成 max_tokens + +
+
+ +
+
+ + + + +""" diff --git a/mlx_streaming/server/app.py b/mlx_streaming/server/app.py new file mode 100644 index 0000000..bb7a8bb --- /dev/null +++ b/mlx_streaming/server/app.py @@ -0,0 +1,104 @@ +"""FastAPI 应用装配:OpenAI 路由 + 引擎调参 API + 极简 admin 页。""" +from __future__ import annotations + +import asyncio +import os +from contextlib import asynccontextmanager + +from fastapi import FastAPI, Request +from fastapi.responses import HTMLResponse + +from mlx_streaming.server.admin import ADMIN_HTML +from mlx_streaming.server.openai import make_openai_router +from mlx_streaming.server.state import EngineManager, rss_bytes +from mlx_streaming.tui.backend import FakeBackend + + +def create_app(args, backend=None) -> FastAPI: + """装配应用。 + + args: argparse Namespace 风格对象(至少含 model/k/max_tokens/expert_slots, + 真后端还需要 expert_dir/mtp_out/qn_config/spec_slots)。 + backend: 显式注入的后端(测试用);None 时按 SPARKLE_FAKE=1 选 FakeBackend, + 否则 MLXBackend。 + """ + if backend is None: + if os.environ.get("SPARKLE_FAKE") == "1": + backend = FakeBackend() + else: + from mlx_streaming.tui.backend import MLXBackend + + backend = MLXBackend(args) + mgr = EngineManager(args, backend) + + @asynccontextmanager + async def lifespan(app): + # 加载放线程且不等待:真模型以分钟计,不能阻塞 uvicorn 启动; + # 就绪前 loaded=False,chat 请求返回 503。 + loop = asyncio.get_running_loop() + app.state.load_future = loop.run_in_executor(None, mgr.load_initial) + yield + + app = FastAPI(title="sparkle server", lifespan=lifespan) + # 本机工具服务:放开 CORS,让浏览器能直连 /v1/models、/api/* 做探测 + from fastapi.middleware.cors import CORSMiddleware + + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["*"], + allow_headers=["*"], + ) + app.state.mgr = mgr + app.include_router(make_openai_router(mgr)) + + @app.get("/api/stats") + async def api_stats(): + return {"tok_per_s": mgr.last_tok_per_s, + "rss_bytes": rss_bytes(), + "expert_slots": mgr.expert_slots, + "loaded": mgr.loaded, + # 专家池「分不到槽被迫落 0 号槽」的累计数。>0 = 有前向拿错专家权重算过, + # 输出不可信(调小 PREFILL_CHUNK 或调大 expert_slots)。健康值恒为 0。 + "unplaced_experts": mgr.unplaced_experts(), + # 专家池命中/读盘指标:hit_rate 低就是慢的直接原因(每 miss 一个要读 ~3.3MB) + "pool": mgr.pool_metrics(), + "model": str(mgr.args.model)} + + @app.get("/api/engine/config") + async def get_engine_config(): + return {"expert_slots": mgr.expert_slots, + "k": mgr.k, + "max_tokens": mgr.max_tokens} + + @app.post("/api/engine/config") + async def set_engine_config(req: Request): + """有多少改多少:k/max_tokens 即时生效;expert_slots 变化触发后台重建。""" + body = await req.json() + if "k" in body: + mgr.k = int(body["k"]) + if "max_tokens" in body: + mgr.max_tokens = int(body["max_tokens"]) + reloading = False + if "expert_slots" in body: + slots = int(body["expert_slots"]) + if slots != mgr.expert_slots: + mgr.expert_slots = slots + mgr.reloading = True # 即刻对 chat 可见(503),不等后台任务起跑 + reloading = True + + async def _bg_rebuild(): + async with mgr.lock: # 等当前生成结束再换引擎 + await asyncio.get_running_loop().run_in_executor( + None, mgr.rebuild) + + asyncio.create_task(_bg_rebuild()) + return {"ok": True, "reloading": reloading, + "expert_slots": mgr.expert_slots, + "k": mgr.k, "max_tokens": mgr.max_tokens} + + @app.get("/admin", response_class=HTMLResponse) + async def admin_page(): + return ADMIN_HTML + + return app diff --git a/mlx_streaming/server/openai.py b/mlx_streaming/server/openai.py new file mode 100644 index 0000000..b2eb3ea --- /dev/null +++ b/mlx_streaming/server/openai.py @@ -0,0 +1,338 @@ +"""OpenAI 兼容端点:/v1/models 与 /v1/chat/completions(SSE 流式 + 普通 JSON)。 + +生成是阻塞调用,放线程池执行;流式用 asyncio.Queue 把后端回调桥接成 SSE 生成器。 +客户端断开时通过 threading.Event 让 backend 的 on_text 返回 True 中断生成。 +""" +from __future__ import annotations + +import asyncio +import json +import os +import re +import threading +import time +import uuid + +from fastapi import APIRouter, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from mlx_streaming import config +from mlx_streaming.server.state import EngineManager + + +def _sse(obj) -> str: + return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n" + + +def _delta_of(prev: str, cur: str) -> str: + """后端上报的是累计全文,这里取出本次新增。 + + 按真实公共前缀切,而不是 `cur[len(prev):]`:后者假设 prev 一定是 cur 的前缀,一旦后端 + 改写了已发出的尾部(历史上流式解码到半个汉字会先发 U+FFFD、下一步再换成真字符), + 长度相同就切出空串,那个字符永久丢失。此处宁可重发一小段,也不丢字。 + """ + if cur.startswith(prev): + return cur[len(prev):] + n = 0 + for a, b in zip(prev, cur): + if a != b: + break + n += 1 + return cur[n:] + + +_MASK_RE = re.compile(r"\[(NAME|ID|TEAM|STATION|CAR|TRAIN):(emp_[a-z0-9]+)\]", re.I) + + +def _mask_notice(content: str) -> str: + """工具结果里带脱敏占位符时,就地追加一句原样透传的约束。 + + 实测(2026-08-02):工具返回 `[TEAM:emp_ypsl004]`,模型答"乘务三组"——把占位符替换成 + 了自己编的名字,用户拿到错数据。系统提示里已有同样的规则,但它在 4000 token 之前, + 而占位符就在眼前;贴着数据再说一次显著更硬,且模型无关。 + 只在真的检出占位符时加,不给普通工具结果增加 token。 + """ + if not content or not _MASK_RE.search(content): + return content + return (f"{content}\n\n[脱敏提示] 上面形如 [TEAM:emp_xxx]、[NAME:emp_xxx]、[ID:emp_xxx] 的是" + "脱敏占位符,界面会自动还原成真实值。引用时**必须连方括号整体原样复制**," + "严禁替换成「乘务一组」这类具体名称(你看不到真实值,写具体名称一定是编的)," + "也不要拆掉方括号或向用户解释占位符。") + + +def _normalize_messages(raw) -> list[dict]: + """OpenAI messages → chat template 可用消息。 + + 关键:assistant 的 tool_calls 与 role=tool 的消息**原样透传**——Qwen chat template + 会把它们渲染成模型训练时学过的格式( 与 包裹)。 + 此前把工具返回改写成普通 user 文本,模型不认为是权威工具结果,会无视真实数据 + 自行编造报告(2026-07 实测复现)。只做 content 分段列表的拍平。 + """ + out = [] + for m in raw: + role = m.get("role", "user") + content = m.get("content") + if isinstance(content, list): # OpenAI 分段 content:[{"type":"text","text":...}] + content = "".join( + p.get("text", "") for p in content + if isinstance(p, dict) and p.get("type") == "text") + if role == "tool": + # 原生透传:模板渲染为 ,模型按训练格式 ground 到工具数据 + item = {"role": "tool", "content": _mask_notice(content or "")} + if m.get("tool_call_id"): + item["tool_call_id"] = m["tool_call_id"] + if m.get("name"): + item["name"] = m["name"] + out.append(item) + continue + if role == "assistant" and m.get("tool_calls"): + # 原生透传:模板渲染为 。arguments 必须是 JSON 字符串(OpenAI 规范) + tcs = [] + for tc in m["tool_calls"]: + fn = tc.get("function", {}) + args = fn.get("arguments") + if args is not None and not isinstance(args, str): + args = json.dumps(args, ensure_ascii=False) + tcs.append({"id": tc.get("id"), "type": "function", + "function": {"name": fn.get("name"), "arguments": args}}) + out.append({"role": "assistant", "content": content or None, + "tool_calls": tcs}) + continue + out.append({"role": role, "content": content or ""}) + return out + + +def _tool_name(t) -> str: + if not isinstance(t, dict): + return "" + fn = t.get("function") + if isinstance(fn, dict) and fn.get("name"): + return str(fn["name"]) + return str(t.get("name") or "") + + +_tools_filter_seen: set = set() + + +def _filter_tools(tools): + """按 config.tools_allow() 白名单裁剪工具,全被裁掉时返回 None(本轮不暴露工具)。 + + 见 config.tools_allow 的说明:这不只是省 token,主要是把 prompt 头部钉稳,让前缀快照 + 能命中——调用方的工具发现是动态的,工具集一变头部就变,每个新会话都要全量 prefill。 + 首次遇到某个「保留/丢弃」组合时打一行日志,便于确认线上到底发来了什么。 + """ + allow = config.tools_allow() + if not allow or not tools: + return tools + allow_set = set(allow) + kept, dropped = [], [] + for t in tools: + (kept if _tool_name(t) in allow_set else dropped).append(t) + if dropped: + sig = (tuple(_tool_name(t) for t in kept), tuple(_tool_name(t) for t in dropped)) + if sig not in _tools_filter_seen: + _tools_filter_seen.add(sig) + print(f"[TOOLS_ALLOW] 保留 {len(kept)} 个 {list(sig[0])}," + f"丢弃 {len(dropped)} 个 {list(sig[1])}。" + f"目的是让 prompt 头部固定、前缀快照命中;设 TOOLS_ALLOW= 可关闭。", + flush=True) + return kept or None + + +def make_openai_router(mgr: EngineManager) -> APIRouter: + router = APIRouter() + + def _model_id() -> str: + return os.path.basename(str(mgr.args.model).rstrip("/")) or "sparkle" + + def _auth(req: Request): + """Bearer 任意非空 token 即放行;设了 SPARKLE_API_KEY 则必须匹配。""" + auth = req.headers.get("authorization", "") + tok = auth[7:].strip() if auth.lower().startswith("bearer ") else "" + want = os.environ.get("SPARKLE_API_KEY", "") + ok = (tok == want) if want else bool(tok) + if ok: + return None + return JSONResponse(status_code=401, content={ + "error": {"message": "unauthorized", "type": "authentication_error"}}) + + def _unavailable(): + msg = "engine reloading" if mgr.reloading else "engine not loaded" + return JSONResponse(status_code=503, content={"error": msg}) + + def _prompt_tokens(messages, tools=None) -> int: + """prompt token 数:真后端用 tokenizer 精确编码;fake 退化为字符数估算。 + + 必须带上 tools:工具定义是由 chat template 渲染进 prompt 的,实测 4 个工具就占 + 1349 token(占总 prompt 三分之一)。漏掉它上报的 usage 会显著偏小,而调用方 + (中台)按 usage 判断要不要压缩上下文,少算就会算漏。 + """ + tok = getattr(mgr.backend, "_tok", None) + if tok is not None: + try: + from mlx_streaming.cli import _encode_chat + + return len(_encode_chat(tok, messages, tools=tools)) + except Exception: # noqa: BLE001 估算失败不阻塞生成 + pass + return sum(len(m.get("content") or "") for m in messages) + + def _usage(prompt_tokens: int, completion_tokens: int) -> dict: + return {"prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens} + + def _chunk(model_id, cmpl_id, created, delta=None, role=False, + finish=None, usage=None): + if usage is not None: # OpenAI 风格:usage chunk 的 choices 为空数组 + return {"id": cmpl_id, "object": "chat.completion.chunk", + "created": created, "model": model_id, + "choices": [], "usage": usage} + d = {} + if role: + d["role"] = "assistant" + if delta: + d["content"] = delta + return {"id": cmpl_id, "object": "chat.completion.chunk", + "created": created, "model": model_id, + "choices": [{"index": 0, "delta": d, "finish_reason": finish}]} + + @router.get("/v1/models") + async def list_models(req: Request): + err = _auth(req) + if err is not None: + return err + return {"object": "list", "data": [{ + "id": _model_id(), "object": "model", + "created": int(time.time()), "owned_by": "sparkle"}]} + + @router.post("/v1/chat/completions") + async def chat_completions(req: Request): + err = _auth(req) + if err is not None: + return err + if not mgr.loaded or mgr.reloading: + return _unavailable() + body = await req.json() + # chat_template_kwargs / enable_thinking:接受并忽略。 + # tools:渲染进 chat template(Qwen 模板把工具定义注入提示词),模型据此输出 + # 文本,客户端兜底解析——server 侧不做原生 tool parser。 + raw_tools = body.get("tools") or None + tools = _filter_tools(raw_tools) + # 「模型没调工具就编数据」有两种成因:调用方压根没发工具,或发了而模型无视。 + # 两者在 server 侧看不出区别(都只是一次普通请求),故每轮记一行收到的工具名。 + print(f"[REQ] 收到工具 {len(raw_tools or [])} 个 " + f"{[_tool_name(t) for t in (raw_tools or [])]} → 渲染 " + f"{len(tools or [])} 个 {[_tool_name(t) for t in (tools or [])]}", flush=True) + # 采样策略(关键):模型 generation_config 要求 do_sample(temp 0.7/top_p 0.8/top_k 20), + # 贪心会诱发幻觉。默认走采样模式(质量优先,非投机);仅显式 temperature=0 走贪心+MTP。 + req_temp = body.get("temperature") + sampling = None + if req_temp is None or float(req_temp) > 0: + gc = getattr(mgr.backend, "_gen_cfg", None) or {} + sampling = { + "temperature": float(req_temp) if req_temp is not None + else float(gc.get("temperature", 0.7)), + "top_p": body.get("top_p") or gc.get("top_p", 0.8), + "top_k": body.get("top_k") or gc.get("top_k", 20), + } + messages = _normalize_messages(body.get("messages") or []) + stream = bool(body.get("stream")) + include_usage = bool((body.get("stream_options") or {}).get("include_usage")) + req_max = body.get("max_tokens") or body.get("max_completion_tokens") + model_id = _model_id() + cmpl_id = "chatcmpl-" + uuid.uuid4().hex[:24] + created = int(time.time()) + loop = asyncio.get_running_loop() + disconnect = threading.Event() # 客户端断开 → on_text 返回 True 中断生成 + + async def _watch_disconnect(): + while not disconnect.is_set(): + if await req.is_disconnected(): + disconnect.set() + return + await asyncio.sleep(0.2) + + if not stream: + async with mgr.lock: # 引擎单对话:并发请求排队等锁 + if not mgr.loaded or mgr.reloading: # 等锁期间可能开始重建 + return _unavailable() + mgr.apply_gen_params(req_max) + prompt_tokens = _prompt_tokens(messages, tools) + watcher = asyncio.create_task(_watch_disconnect()) # 客户端断开也中断,防僵尸生成堵锁 + try: + res = await loop.run_in_executor( + None, lambda: mgr.backend.generate( + messages, lambda _f, _n: disconnect.is_set(), + tools=tools, sampling=sampling)) + except Exception as e: # noqa: BLE001 + return JSONResponse(status_code=500, content={"error": str(e)}) + finally: + disconnect.set() + watcher.cancel() + mgr.last_tok_per_s = res.tok_per_s + return {"id": cmpl_id, "object": "chat.completion", + "created": created, "model": model_id, + "choices": [{"index": 0, "finish_reason": "stop", + "message": {"role": "assistant", + "content": res.text}}], + "usage": _usage(prompt_tokens, res.n_tokens)} + + async def event_stream(): + async with mgr.lock: + if not mgr.loaded or mgr.reloading: + yield _sse({"error": "engine reloading"}) + yield "data: [DONE]\n\n" + return + mgr.apply_gen_params(req_max) + prompt_tokens = _prompt_tokens(messages, tools) + q = asyncio.Queue() + + def on_text(full, _n): + if disconnect.is_set(): + return True + loop.call_soon_threadsafe(q.put_nowait, full) + return False + + def _run(): # 阻塞生成,放线程池;结束/异常都经由 queue 通知 + try: + res = mgr.backend.generate(messages, on_text, tools=tools, + sampling=sampling) + loop.call_soon_threadsafe(q.put_nowait, ("__done__", res)) + except Exception as e: # noqa: BLE001 + loop.call_soon_threadsafe(q.put_nowait, ("__error__", e)) + + watcher = asyncio.create_task(_watch_disconnect()) + loop.run_in_executor(None, _run) + prev = "" + result = None + yield _sse(_chunk(model_id, cmpl_id, created, role=True)) + try: + while True: + item = await q.get() + if isinstance(item, tuple): + if item[0] == "__done__": + result = item[1] + break + raise item[1] + delta = _delta_of(prev, item) + prev = item + if delta: + yield _sse(_chunk(model_id, cmpl_id, created, + delta=delta)) + finally: + disconnect.set() # 客户端断开/生成结束都确保后端被中断 + watcher.cancel() + if result is not None: + mgr.last_tok_per_s = result.tok_per_s + yield _sse(_chunk(model_id, cmpl_id, created, finish="stop")) + if include_usage: + n = result.n_tokens if result is not None else 0 + yield _sse(_chunk(model_id, cmpl_id, created, + usage=_usage(prompt_tokens, n))) + yield "data: [DONE]\n\n" + + return StreamingResponse(event_stream(), + media_type="text/event-stream") + + return router diff --git a/mlx_streaming/server/state.py b/mlx_streaming/server/state.py new file mode 100644 index 0000000..288386a --- /dev/null +++ b/mlx_streaming/server/state.py @@ -0,0 +1,101 @@ +"""引擎单例管理:加载状态、生成串行锁、运行期调参与后台重建。""" +from __future__ import annotations + +import asyncio +import resource +import sys + +from mlx_streaming.tui.backend import FakeBackend + + +def rss_bytes() -> int: + """当前进程常驻内存峰值(字节)。macOS 的 ru_maxrss 直接是字节;Linux 是 KB。""" + ru = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return int(ru if sys.platform == "darwin" else ru * 1024) + + +class EngineManager: + """持有一个 ChatBackend 单例;引擎是单对话的,同一时刻只允许一轮生成。 + + asyncio.Lock 做排队:第二个并发请求等锁而非直接失败。k / max_tokens / expert_slots + 是运行期可调参数,由 /api/engine/config 修改;k 与 max_tokens 即时生效(作为后续 + OpenAI 请求的缺省值),expert_slots 变化触发后台重建。 + """ + + def __init__(self, args, backend): + self.args = args # argparse Namespace 风格,重建 MLXBackend 时复用 + self.backend = backend + self.loaded = False + self.reloading = False + self.load_error: str | None = None + self.lock = asyncio.Lock() + self.last_tok_per_s = 0.0 + self.k = getattr(args, "k", 3) + self.max_tokens = getattr(args, "max_tokens", 4096) + self.expert_slots = getattr(args, "expert_slots", 32) + + def load_initial(self, on_status=lambda m: None) -> None: + """启动加载(放线程里跑,真模型以分钟计);就绪前 loaded=False。""" + try: + self.backend.load(on_status) + self.loaded = True + except Exception as e: # noqa: BLE001 加载失败不拖垮 server,状态暴露在 /api/stats + self.load_error = f"{type(e).__name__}: {e}" + + def rebuild(self, on_status=lambda m: None) -> None: + """后台重建引擎(fake 模式只换实例;真后端按更新后的 args 重新装配 + 预热)。""" + self.reloading = True + self.loaded = False + try: + self.args.expert_slots = self.expert_slots + old = self.backend + if isinstance(old, FakeBackend): + self.backend = FakeBackend(reply=old.reply) + else: + from mlx_streaming.tui.backend import MLXBackend + + self.backend = MLXBackend(self.args) + self.backend.load(on_status) + self.loaded = True + except Exception as e: # noqa: BLE001 + self.load_error = f"{type(e).__name__}: {e}" + finally: + self.reloading = False + + def _resident_pool(self): + """拿常驻专家池(fake 后端 / 非双源路径下为 None)。""" + store = getattr(getattr(self.backend, "_model", None), "_expert_store", None) + return getattr(store, "_resident", None) + + def unplaced_experts(self) -> int: + """专家池累计「分不到槽落 0 号槽」数(见 ResidentExpertPool.unplaced)。 + + 取不到(fake 后端 / 非双源路径)返回 0。这是数值正确性的看门指标:非 0 就说明 + 某些前向是拿错专家权重算的,回答不可信。 + """ + return int(getattr(self._resident_pool(), "unplaced", 0) or 0) + + def pool_metrics(self) -> dict: + """专家池命中/读盘指标：定位「慢」到底是池不够大还是走了慢路径。 + + decode 每步要取 层数×top_k 个专家，miss 一个就得读一份专家权重(8bit ≈3.3MB)， + 所以 hit_rate 直接决定速度。gpu_fallback 记的是本该零 host 往返的快路径退回读盘的次数。 + """ + rp = self._resident_pool() + if rp is None: + return {} + hits, misses = int(getattr(rp, "hits", 0)), int(getattr(rp, "misses", 0)) + return {"hits": hits, "misses": misses, + "hit_rate": round(hits / (hits + misses), 4) if hits + misses else 0.0, + "gpu_fastpath": int(getattr(rp, "gpu_fastpath", 0)), + "gpu_fallback": int(getattr(rp, "gpu_fallback", 0)), + "prefetch_hits": int(getattr(rp, "prefetch_hits", 0)), + "prefetch_loads": int(getattr(rp, "prefetch_loads", 0))} + + def apply_gen_params(self, req_max_tokens=None) -> int: + """把当前调参写入后端 args(单飞前提下安全),返回本轮有效 max_tokens。""" + eff = req_max_tokens or self.max_tokens + if hasattr(self.backend, "args"): + self.backend.args.k = self.k + self.backend.args.max_tokens = eff + return eff diff --git a/mlx_streaming/tui/__init__.py b/mlx_streaming/tui/__init__.py new file mode 100644 index 0000000..731c45d --- /dev/null +++ b/mlx_streaming/tui/__init__.py @@ -0,0 +1,9 @@ +"""sparkle 全屏 TUI 包。run_tui 延迟 import app(避免无谓加载 textual)。""" + + +def run_tui(backend, args) -> int: + """启动全屏 TUI。backend 为 ChatBackend 实现,args 为解析后的命令行参数。""" + from mlx_streaming.tui.app import SparkleApp + + SparkleApp(backend, args).run() + return 0 diff --git a/mlx_streaming/tui/app.py b/mlx_streaming/tui/app.py new file mode 100644 index 0000000..4fff697 --- /dev/null +++ b/mlx_streaming/tui/app.py @@ -0,0 +1,334 @@ +"""sparkle 全屏 TUI(Textual 实现,视觉对标 opencode)。 + +界面只依赖 ChatBackend 接口。加载与每轮生成均在 worker 线程执行, +通过 call_from_thread 把状态/流式文本安全回推到 UI 线程。 +""" +from __future__ import annotations + +import os +import time +from collections import deque +from typing import Optional + +from rich.console import Group +from rich.markdown import Markdown as RichMarkdown +from rich.text import Text +from textual import work +from textual.app import App, ComposeResult +from textual.containers import VerticalScroll +from textual.widgets import Input, Static + +from mlx_streaming.core.mem import snapshot +from mlx_streaming.tui.backend import ChatBackend, GenResult +from mlx_streaming.tui.banner import LOGO + +_ACCENT = "#2dd4bf" + + +def _fmt_gb(nbytes: int) -> str: + return f"{nbytes / 1e9:.2f} GB" + + +def _mem_suffix(peak: bool = False) -> str: + """状态栏内存后缀。显存外置流式 MoE 的核心卖点是低内存,故常驻展示活跃占用; + 结束时附带峰值。取 MLX 统一内存计数器,开销极低。取不到(0)时返回空串,不干扰状态栏。""" + snap = snapshot() + if snap.mlx_active_bytes <= 0: + return "" + s = f" · 内存 {_fmt_gb(snap.mlx_active_bytes)}" + if peak and snap.mlx_peak_bytes > 0: + s += f"(峰值 {_fmt_gb(snap.mlx_peak_bytes)})" + return s + +# 生成中状态栏用滑动窗口算「瞬时」tok/s 的时间窗(秒);越小越灵敏、越大越平滑。 +_TPS_WINDOW = 1.0 + +_HELP = ( + "可用命令:\n" + " /help 显示本帮助\n" + " /reset 清空对话历史(保留 system)\n" + " /clear 清空对话区显示\n" + " /exit 退出\n\n" + "快捷键:Enter 发送 · Esc 中断生成 · Ctrl+C 退出" +) + + +def _short_model(path: str) -> str: + # 只取模型路径末段,避免顶栏/状态栏过长 + return os.path.basename(path.rstrip("/")) or path + + +class ChatMessage(Static): + """一条对话消息。role ∈ {'user','assistant'};助手消息支持流式更新与 Markdown 收尾。""" + + def __init__(self, role: str, text: str = "", *, final: bool = True): + super().__init__(classes=f"msg {role}") + self.role = role + self.text = text + self.final = final + self._refresh_content() + + def stream(self, text: str) -> None: + # 流式过程中用纯文本渲染,避免半截 Markdown 抖动 + self.text = text + self.final = False + self._refresh_content() + + def finalize(self, text: str) -> None: + # 收尾时切换为 Markdown 渲染 + self.text = text + self.final = True + self._refresh_content() + + def _refresh_content(self) -> None: + if self.role == "user": + header = Text("› 你", style="bold") + body = Text(self.text) + else: + header = Text("⏺ sparkle", style=f"bold {_ACCENT}") + if not self.text: + body = Text("正在思考…", style="dim italic") + elif self.final: + body = RichMarkdown(self.text) + else: + body = Text(self.text) + self.update(Group(header, body)) + + +class SparkleApp(App): + CSS_PATH = "styles.tcss" + BINDINGS = [ + ("escape", "interrupt", "中断"), + ("ctrl+c", "quit", "退出"), + ] + + def __init__(self, backend: ChatBackend, args): + super().__init__() + self.backend = backend + self.args = args + self._messages: list[dict] = [] + # 保留 system 消息,/reset 时不清空这部分 + if getattr(args, "system", None): + self._messages.append({"role": "system", "content": args.system}) + self._base_len = len(self._messages) + self._busy = False + self._stop = False + self._cur: Optional[ChatMessage] = None + # 首 token 到达时刻与基准 token 数,用于结束时计算「累计解码平均」tok/s + # (排除 prefill)。_gen_t0 为 0.0 表示本轮尚未收到首 token。 + self._gen_t0 = 0.0 + self._n0 = 0 + # 生成中「瞬时」tok/s 的滑动窗口采样:每项为 (时刻, 累计 token 数)。 + self._tps_window = _TPS_WINDOW + self._samples: deque[tuple[float, int]] = deque() + + def compose(self) -> ComposeResult: + yield Static(self._top_text(), id="top") + yield VerticalScroll(id="chat") + # 提示放到边框标题里,不用文本区的 placeholder: + # 长占位符在部分终端增量重绘时不会被擦除,打字后会残留「后面还有字」, + # 直到全量重绘(Enter/截图/resize)才消失;边框标题在边框上,不受此影响。 + yield Input(id="prompt") + yield Static("", id="status") + + def on_mount(self) -> None: + inp = self.query_one("#prompt", Input) + inp.border_title = "输入消息 · Enter 发送 · Esc 中断 · /help" + # 加载完成前禁用输入框,避免用户在模型就绪前发消息 + inp.disabled = True + self._set_status("加载模型中…") + self._load() + + def on_input_changed(self, event: Input.Changed) -> None: + # 兜底:强制整屏重绘,清除个别终端增量重绘遗留的输入残影 + # (等价于截图/resize 触发的全量刷新)。 + self.refresh() + + def _top_text(self) -> Text: + t = Text() + t.append(LOGO, style=f"bold {_ACCENT}") + t.append(f" {_short_model(self.args.model)} · k={self.args.k} · " + f"{self.args.max_tokens} tok", style="dim") + return t + + def _set_status(self, text: str) -> None: + self.query_one("#status", Static).update( + f" {_short_model(self.args.model)} · {text}") + + @work(thread=True, exclusive=True, group="load") + def _load(self) -> None: + # 在 worker 线程执行阻塞加载,回调需经 call_from_thread 回到 UI 线程 + def on_status(msg: str): + # 应用已退出时不再回推,避免 call_from_thread 抛错 + if self.is_running: + self.call_from_thread(self._set_status, f"加载中 · {msg}") + + try: + self.backend.load(on_status) + except Exception as e: # noqa: BLE001 + if self.is_running: + self.call_from_thread(self._on_load_failed, str(e)) + return + if self.is_running: + self.call_from_thread(self._on_load_done) + + def _enable_input(self) -> None: + """重新启用并聚焦输入框(加载完成/一轮生成结束后统一调用)。 + + 刚把 disabled 置 False 时 can_focus 状态可能还没刷新,直接 focus 偶发无效 + (表现为输入框失焦时占位提示不消失);故延到下一次刷新后再聚焦,确保稳定拿到焦点。 + """ + inp = self.query_one("#prompt", Input) + inp.disabled = False + self.call_after_refresh(inp.focus) + + def _on_load_done(self) -> None: + self._set_status(f"就绪{_mem_suffix()}") + self._enable_input() + + def _on_load_failed(self, err: str) -> None: + self._add("assistant", f"模型加载失败:{err}\n\n请检查路径后用 /exit 退出重试。") + self._set_status("加载失败") + + def on_input_submitted(self, event: Input.Submitted) -> None: + text = event.value.strip() + event.input.value = "" + if not text: + return + if text.startswith("/"): + self._command(text) + return + if self._busy: + return + self._start(text) + + def _command(self, cmd: str) -> None: + if cmd in ("/exit", "/quit"): + self.exit() + elif cmd == "/help": + self._add("assistant", _HELP) + elif cmd in ("/reset", "/clear") and self._busy: + # 生成中改动历史/移除正在流式的 _cur 会让 _on_stream/_on_done 操作已卸载组件,故拒绝 + self._add("assistant", "生成中,请等本轮结束或按 Esc 中断后再执行该命令。") + elif cmd == "/reset": + # 只清空 system 之后的历史 + del self._messages[self._base_len:] + self._add("assistant", "对话历史已清空。") + elif cmd == "/clear": + self.query_one("#chat", VerticalScroll).remove_children() + else: + self._add("assistant", f"未知命令:{cmd}(/help 查看可用命令)") + + def _start(self, user_text: str) -> None: + self._messages.append({"role": "user", "content": user_text}) + self._add("user", user_text) + self._cur = self._add("assistant", "", final=False) + self._busy = True + self._stop = False + # 置 0 表示还没收到首 token;真正的计时起点推迟到第一次 _on_stream。 + self._gen_t0 = 0.0 + self._n0 = 0 + self._samples.clear() + self.query_one("#prompt", Input).disabled = True + self._set_status("思考中…") + self._generate(list(self._messages)) + + @work(thread=True, exclusive=True, group="gen") + def _generate(self, messages: list[dict]) -> None: + # 生成在 worker 线程运行;on_text 返回 self._stop 让后端可提前中断 + def on_text(full: str, n_tokens: int) -> bool: + # 应用已退出时,call_from_thread 会抛错;此时直接请求停止,避免 worker 线程未捕获异常。 + if not self.is_running: + return True + try: + self.call_from_thread(self._on_stream, full, n_tokens) + except Exception: # noqa: BLE001 + return True + return self._stop + + try: + result = self.backend.generate(messages, on_text) + except Exception as e: # noqa: BLE001 + if self.is_running: + self.call_from_thread(self._on_error, str(e)) + return + if self.is_running: + self.call_from_thread(self._on_done, result) + + def _record_sample(self, now: float, n_tokens: int) -> None: + """记录一次采样,并丢弃早于滑动窗口的旧点(保留跨越窗口边界的那一个)。""" + self._samples.append((now, n_tokens)) + w = self._tps_window + # 当第二个点仍比窗口还老时,第一个点已冗余,可丢弃; + # 循环后 _samples[0] 恰好是刚跨过窗口边界的采样,窗口长度 ≈ w。 + while len(self._samples) >= 2 and now - self._samples[1][0] >= w: + self._samples.popleft() + + def _window_tps(self, now: float) -> Optional[float]: + """按滑动窗口算瞬时 tok/s;样本不足(不够两点或时间差为 0)时返回 None。""" + if len(self._samples) < 2: + return None + t0, n0 = self._samples[0] + dt = now - t0 + if dt <= 0: + return None + return (self._samples[-1][1] - n0) / dt + + def _on_stream(self, full: str, n_tokens: int) -> None: + if self._cur is not None: + self._cur.stream(full) + self._scroll_end() + now = time.monotonic() + # 首次回调:prefill 刚结束,记下计时起点与基准 token 数,供结束时算累计解码平均。 + if self._gen_t0 == 0.0: + self._gen_t0 = now + self._n0 = n_tokens + # 生成中显示滑动窗口「瞬时」速度:一直在动,能反映后期变慢,不被历史平均拖住。 + self._record_sample(now, n_tokens) + tps = self._window_tps(now) + mem = _mem_suffix() + if tps is None: + self._set_status(f"思考中 · {n_tokens} tok{mem}") + else: + self._set_status(f"思考中 · {n_tokens} tok · {tps:.1f} tok/s{mem}") + + def _on_done(self, result: GenResult) -> None: + if self._cur is not None: + self._cur.finalize(result.text) + self._messages.append({"role": "assistant", "content": result.text}) + self._busy = False + self._cur = None + suffix = " · 已中断" if result.stopped else "" + # 与流式状态栏同口径:从首 token 起、按解码 token 数算,避免结束瞬间数字回落。 + # 若本轮没触发过流式(_gen_t0 仍为 0),退回后端上报的 tok/s。 + dt = time.monotonic() - self._gen_t0 + if self._gen_t0 > 0.0 and dt > 0: + tps = (result.n_tokens - self._n0) / dt + else: + tps = result.tok_per_s + self._set_status( + f"就绪 · {result.n_tokens} tok · {tps:.1f} tok/s{_mem_suffix(peak=True)}{suffix}") + self._enable_input() + + def _on_error(self, err: str) -> None: + if self._cur is not None: + self._cur.finalize(f"生成出错:{err}") + self._busy = False + self._cur = None + self._set_status("就绪(上一轮出错)") + self._enable_input() + + def action_interrupt(self) -> None: + # 仅在生成中时置位中断标志,由 on_text 闭包读取 + if self._busy: + self._stop = True + self._set_status("正在中断…") + + def _add(self, role: str, text: str, *, final: bool = True) -> ChatMessage: + msg = ChatMessage(role, text, final=final) + self.query_one("#chat", VerticalScroll).mount(msg) + self._scroll_end() + return msg + + def _scroll_end(self) -> None: + self.query_one("#chat", VerticalScroll).scroll_end(animate=False) diff --git a/mlx_streaming/tui/backend.py b/mlx_streaming/tui/backend.py new file mode 100644 index 0000000..c7a6471 --- /dev/null +++ b/mlx_streaming/tui/backend.py @@ -0,0 +1,436 @@ +"""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"]) diff --git a/mlx_streaming/tui/banner.py b/mlx_streaming/tui/banner.py new file mode 100644 index 0000000..0f85b06 --- /dev/null +++ b/mlx_streaming/tui/banner.py @@ -0,0 +1,3 @@ +"""sparkle 顶栏 logo。保持单行,便于与模型信息拼在同一行。""" + +LOGO = "▚ sparkle" diff --git a/mlx_streaming/tui/styles.tcss b/mlx_streaming/tui/styles.tcss new file mode 100644 index 0000000..de1168a --- /dev/null +++ b/mlx_streaming/tui/styles.tcss @@ -0,0 +1,41 @@ +Screen { + background: #0b0e14; +} + +#top { + height: auto; + padding: 1 2; + background: #11151f; + color: #e6e6e6; + border-bottom: solid #2a2f3a; +} + +#chat { + padding: 1 2; + height: 1fr; +} + +.msg { + margin: 1 0; + padding: 0 1; + height: auto; +} + +.msg.assistant { + border-left: solid #2dd4bf; + padding-left: 2; +} + +#prompt { + margin: 0 2; + border: round #2dd4bf; + border-title-color: #8b93a7; + background: #11151f; +} + +#status { + height: 1; + padding: 0 2; + background: #11151f; + color: #8b93a7; +} diff --git a/models/README.md b/models/README.md new file mode 100644 index 0000000..56063a8 --- /dev/null +++ b/models/README.md @@ -0,0 +1,49 @@ +# 模型与专家数据（不入库，请本机下载/生成） + +## 为什么要做 + +目标不是「把 80B 常驻塞进 64G」，而是用本机 **64G 统一内存** 做压力验证：把约 **77G 专家权重** 以流式专家池压到约 **33G 常驻**，在 8bit 下跑通 Qwen3-Next-80B-A3B。 + +验证点对应下游的 **黄金拐点**：Qwen3 **30B 8bit A3B** 一类模型，用同一套流式装载思路，能放进 **32G / 16G** 统一内存机器，相对「整模进内存」可 **节省 50% 以上统一内存**。本仓库的 80B @ 64G 是这条路径的上限探针；量产形态落在更小机型上的 30B 拐点模型。 + +--- + +## 运行依赖版本 + +权重本身不绑版本，但本仓库预编译扩展与锁文件对齐如下（换版本可能无法 `import native_moe_ext` 或加载 8bit）： + +| 项 | 要求 / 已验证 | +|----|----------------| +| Python | **3.14.***（已验证 3.14.6；须匹配 `native_moe_ext.cpython-314-darwin.so`） | +| mlx | `>=0.31`（锁文件 **0.31.2**） | +| mlx-lm | `>=0.31`（锁文件 **0.31.3**） | +| numpy | `>=2.0`（锁文件 **2.4.6**） | +| fastapi | `>=0.115`（锁文件 **0.140.0**） | +| uvicorn | `>=0.30`（锁文件 **0.51.0**） | +| textual | `>=0.80`（锁文件 **8.2.8**） | + +安装：在仓库根目录 `uv sync`（或使用已有 `.venv`）。完整约束见根目录 `pyproject.toml` / `uv.lock`。换 Python 小版本需 `make -C native/ext native_moe_ext` 重编 `.so`。 + +权重体积约 **160GB**，不进 Git。按下面准备后目录应类似： + +``` +models/ +├── Qwen3-Next-80B-A3B-Instruct-MLX-8bit/ # ~79GB 主模型 +├── experts_8bit_g64/ # ~77GB 专家 blob（运行时常驻约 33G） +└── qn_mtp_weights.safetensors # ~3.1GB MTP +``` + +## 下载（NAS） + +全部模型与专家数据从 NAS 共享下载： + +- 链接：https://ug.link/dxp4800plus-8761/filemgr/share-download/?id=86b004b90f4a477e84c7f65c87943b2f +- 密码：`sparkle` + +解压/放置到本仓库 `models/` 下，目录结构与上文一致即可。 + +> 说明：本仓库当前跑通的是 **80B** 探针；包内若含 **Qwen3-30B-A3B-Instruct-2507-MLX-8bit**，供下游 16G/32G 机型对照，不替代本机 80B 数据准备。 + +## 可选 + +`prefix_snapshots/` 会在首次跑通后自动生成，可不预置。 diff --git a/native/ext/CMakeLists.txt b/native/ext/CMakeLists.txt new file mode 100644 index 0000000..2df3844 --- /dev/null +++ b/native/ext/CMakeLists.txt @@ -0,0 +1,32 @@ +cmake_minimum_required(VERSION 3.22) +project(native_moe_mlx_ext LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) +find_package(nanobind CONFIG REQUIRED) +find_package(MLX CONFIG REQUIRED) + +nanobind_add_module( + native_moe_ext + NB_STATIC + STABLE_ABI + LTO + NOMINSIZE + NB_DOMAIN mlx + compute/fused_moe.cpp + io/blob_load.cpp + io/bg_reader.cpp + prefetch/prefetch.cpp + pool/owned_pool.cpp + pool/side_region.cpp + pool/demand.cpp + bindings.cpp +) + +target_link_libraries(native_moe_ext PRIVATE mlx) + +if(BUILD_SHARED_LIBS) + target_link_options(native_moe_ext PRIVATE -Wl,-rpath,@loader_path) +endif() diff --git a/native/ext/Makefile b/native/ext/Makefile new file mode 100644 index 0000000..9ecea0e --- /dev/null +++ b/native/ext/Makefile @@ -0,0 +1,14 @@ +# 生产 MLX 扩展构建：cmake 把 native_moe_mlx_ext.cpp 编成 +# mlx_streaming/native_moe_ext$(EXT_SUFFIX).so（nanobind + MLX Primitive）。 +PYTHON ?= ../../.venv/bin/python +PY_SITE := $(abspath ../../.venv/lib/python3.12/site-packages) +NANOBIND_DIR := $(shell $(PYTHON) -m nanobind --cmake_dir) +MLX_CMAKE_DIR := $(PY_SITE)/mlx/share/cmake/MLX + +native_moe_ext: + cmake -S . -B build -DPython_EXECUTABLE=$(abspath $(PYTHON)) -Dnanobind_DIR=$(NANOBIND_DIR) -DMLX_DIR=$(MLX_CMAKE_DIR) -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$(abspath ../../mlx_streaming) + cmake --build build --target native_moe_ext -j + +clean: + rm -rf build + rm -f ../../mlx_streaming/native_moe_ext*.so diff --git a/native/ext/bindings.cpp b/native/ext/bindings.cpp new file mode 100644 index 0000000..99237a9 --- /dev/null +++ b/native/ext/bindings.cpp @@ -0,0 +1,153 @@ +#include "compute/fused_moe.h" +#include "io/blob_load.h" +#include "io/bg_reader.h" +#include "prefetch/prefetch.h" +#include "pool/owned_pool.h" +#include "pool/side_region.h" +#include "pool/demand.h" +using namespace nb::literals; + +NB_MODULE(native_moe_ext, m) { + nb::module_::import_("mlx.core"); + m.doc() = "Native MLX MoE extension."; + + // ---- 融合 MoE 计算核（native_fused.cpp）---- + m.def( + "fused_moe", + &fused_moe, + "x"_a, + "expert_ids"_a, + "scores"_a, + "compute_dir"_a, + "layer"_a, + "hidden"_a, + "inter"_a, + "group"_a, + "bits"_a, + "num_experts"_a, + "synthetic"_a, + nb::kw_only(), + "stream"_a = nb::none()); + m.def( + "fused_moe_staged", + &fused_moe_staged, + "x"_a, + "scores"_a, + "gate_w"_a, + "gate_s"_a, + "gate_b"_a, + "up_w"_a, + "up_s"_a, + "up_b"_a, + "down_w"_a, + "down_s"_a, + "down_b"_a, + "hidden"_a, + "inter"_a, + "group"_a, + "bits"_a, + nb::kw_only(), + "stream"_a = nb::none()); + // ---- [1] blob 直读 ---- + m.def( + "blob_load", + &blob_load, + "path"_a, + "expert_ids"_a, + "stride"_a, + nb::kw_only(), + "stream"_a = nb::none()); + // ---- [2] 轻量预取(无 staging，仅预热 page cache) ---- + m.def( + "prefetch_on_complete", + &prefetch_on_complete, + "expert_ids"_a, + "path"_a, + "stride"_a, + "do_read"_a = true, + nb::kw_only(), + "stream"_a = nb::none()); + // ---- [3] staging miss→hit + 完成回调时刻探针（诊断）---- + m.def( + "prefetch_into_staging", + &prefetch_into_staging, + "staging"_a, + "expert_ids"_a, + "layer"_a, + "gen"_a, + "path"_a, + "stride"_a, + "resident"_a, + "cap"_a, + "parallel"_a, + nb::kw_only(), + "stream"_a = nb::none()); + m.def("prefetch_staging_take", &prefetch_staging_take, "layer"_a); + m.def("staging_hprof_enable", &staging_hprof_enable, "on"_a); + m.def("staging_hprof_now", &staging_hprof_now); + m.def("staging_hprof_get", &staging_hprof_get); + // ---- [4] 段散写侧区缓存（zero-copy dual-source 默认路径）---- + m.def("prefetch_pool_sideregion", &prefetch_pool_sideregion, + "pool_list"_a, "seg_nbytes"_a, "expert_ids"_a, "layer"_a, "path"_a, "stride"_a, + "resident"_a, "spec_slots"_a, "base_row"_a, "gen"_a = 0, nb::kw_only(), + "stream"_a = nb::none()); + m.def("sideregion_contents", &sideregion_contents, "layer"_a, "gen"_a = 0); + m.def("sideregion_kv", &sideregion_kv, "layer"_a, "gen"_a = 0); + m.def("sideregion_reset", &sideregion_reset); + m.def("sideregion_drain", &sideregion_drain, + nb::call_guard()); // 等待时释放 GIL + // ---- [5] Route 3 owned 池底座（C++ 拥有 buffer + 直写，删 MLX scatter）---- + m.def("pool_owned_zeros", &pool_owned_zeros, "shape"_a, "dtype"_a); + m.def("pool_write_rows", &pool_write_rows, "pool_list"_a, "srcs_flat"_a, "slots"_a); + m.def("pool_write_stacked", &pool_write_stacked, "pool_list"_a, "stacked_list"_a, "slots"_a); + m.def("array_data_ptr", &array_data_ptr, "a"_a); + // ---- [6] 方案B 真实区槽状态 C++ 接管 + demand_dual ---- + m.def("real_init", &real_init, "layer"_a, "cap"_a); + m.def("real_region_contents", &real_region_contents, "layer"_a); + m.def("real_region_count", &real_region_count, "layer"_a); + m.def("real_reset", &real_reset); + m.def("real_pin", &real_pin, "layer"_a, "experts"_a); + m.def("real_pinned_count", &real_pinned_count, "layer"_a); + m.def("real_freq_dump", &real_freq_dump); + m.def("demand_dual", &demand_dual, "inds"_a, "pool_list"_a, "seg_nbytes"_a, "layer"_a, + "side_gen"_a, "path"_a, "stride"_a, "cap"_a, "lfu"_a, "decay_interval"_a, nb::kw_only(), + "stream"_a = nb::none()); + m.def("demand_last_stats", &demand_last_stats); + m.def("demand_timings", &demand_timings); + m.def("demand_timing_enable", &demand_timing_enable, "on"_a); + m.def("real_debug_place", &real_debug_place, "layer"_a, "experts_flat"_a, "cap"_a, "lfu"_a, + "decay_interval"_a); + // ---- [7] 自由后台读线程 ---- + m.def("bg_reader_start", &bg_reader_start, "workers"_a = 1, "low_cap"_a = 0); + m.def("bg_reader_submit", &bg_reader_submit, + "dst"_a, "experts"_a, "rows"_a, "path"_a, "stride"_a, "ticket"_a, "prio"_a = 0); + m.def("bg_reader_ready", &bg_reader_ready, "ticket"_a); + m.def("bg_reader_wait", &bg_reader_wait, "ticket"_a, + nb::call_guard()); // 阻塞等时释放 GIL + m.def("bg_reader_stop", &bg_reader_stop); + m.def("bg_pread_into_pool", &bg_pread_into_pool, + "dst"_a, "seg_off"_a, "seg_nb"_a, "slot"_a, "expert"_a, + "path"_a, "stride"_a, "ticket"_a, "prio"_a = 0, "nocache"_a = true); + // ---- 融合 MoE 计算核（slots 版，native_fused.cpp）---- + m.def( + "fused_moe_slots", + &fused_moe_slots, + "x"_a, + "local_slots"_a, + "scores"_a, + "gate_w"_a, + "gate_s"_a, + "gate_b"_a, + "up_w"_a, + "up_s"_a, + "up_b"_a, + "down_w"_a, + "down_s"_a, + "down_b"_a, + "hidden"_a, + "inter"_a, + "group"_a, + "bits"_a, + nb::kw_only(), + "stream"_a = nb::none()); +} diff --git a/native/ext/common.h b/native/ext/common.h new file mode 100644 index 0000000..495cd0c --- /dev/null +++ b/native/ext/common.h @@ -0,0 +1,34 @@ +// 共用前导：系统/Metal/MLX/nanobind 头 + 命名空间别名。所有 TU 都包含。 +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mlx/backend/metal/device.h" +#include "mlx/mlx.h" +#include "mlx/primitives.h" + +namespace nb = nanobind; +namespace mx = mlx::core; diff --git a/native/ext/compute/fused_moe.cpp b/native/ext/compute/fused_moe.cpp new file mode 100644 index 0000000..0a94b0d --- /dev/null +++ b/native/ext/compute/fused_moe.cpp @@ -0,0 +1,664 @@ +#include "fused_moe.h" + +struct FusedParams { + int experts; + int hidden; + int inter; + int group_size; + int bits; + int k; +}; + +struct MappedProj { + void* w{MAP_FAILED}; + void* s{MAP_FAILED}; + void* b{MAP_FAILED}; + size_t wn{0}, sn{0}, bn{0}; + int wfd{-1}, sfd{-1}, bfd{-1}; +}; + +struct ProjectionShape { + int out_dim; + int in_dim; + size_t weight_per; + size_t scale_per; +}; + +static size_t file_size(int fd) { + struct stat st {}; + if (fstat(fd, &st) != 0) { + throw std::runtime_error("fstat failed"); + } + return static_cast(st.st_size); +} + +static void* map_file(const std::string& path, size_t& nbytes, int& fd) { + fd = open(path.c_str(), O_RDONLY); + if (fd < 0) { + throw std::runtime_error("failed to open " + path); + } + nbytes = file_size(fd); + void* p = mmap(nullptr, nbytes, PROT_READ, MAP_PRIVATE, fd, 0); + if (p == MAP_FAILED) { + throw std::runtime_error("failed to mmap " + path); + } + return p; +} + +static MappedProj map_proj(const std::string& dir, int layer, const std::string& proj) { + MappedProj m; + std::string layer_name = "layer" + std::string(layer < 10 ? "0" : "") + std::to_string(layer); + std::string base = dir + "/" + layer_name + "." + proj; + m.w = map_file(base + ".weight.bin", m.wn, m.wfd); + m.s = map_file(base + ".scales.bin", m.sn, m.sfd); + m.b = map_file(base + ".biases.bin", m.bn, m.bfd); + return m; +} + +static void unmap_proj(MappedProj& m) { + if (m.w != MAP_FAILED) munmap(m.w, m.wn); + if (m.s != MAP_FAILED) munmap(m.s, m.sn); + if (m.b != MAP_FAILED) munmap(m.b, m.bn); + if (m.wfd >= 0) close(m.wfd); + if (m.sfd >= 0) close(m.sfd); + if (m.bfd >= 0) close(m.bfd); +} + +static ProjectionShape proj_shape(int out_dim, int in_dim, int group, int bits) { + int words = (in_dim * bits) / 32; + int groups = in_dim / group; + return ProjectionShape{ + out_dim, + in_dim, + static_cast(out_dim) * words * sizeof(uint32_t), + static_cast(out_dim) * groups * sizeof(uint16_t), + }; +} + +static uint16_t fp32_to_bf16(float x) { + uint32_t u = 0; + std::memcpy(&u, &x, sizeof(u)); + return static_cast(u >> 16); +} + +static void fill_synthetic( + MTL::Buffer* wbuf, + MTL::Buffer* sbuf, + MTL::Buffer* bbuf, + uint32_t seed) { + std::mt19937 rng(seed); + auto* w = reinterpret_cast(wbuf->contents()); + size_t wn = wbuf->length() / sizeof(uint32_t); + for (size_t i = 0; i < wn; ++i) w[i] = rng(); + auto* s = reinterpret_cast(sbuf->contents()); + auto* b = reinterpret_cast(bbuf->contents()); + size_t sn = sbuf->length() / sizeof(uint16_t); + uint16_t scale = fp32_to_bf16(0.001f); + uint16_t bias = fp32_to_bf16(-0.032f); + for (size_t i = 0; i < sn; ++i) { + s[i] = scale; + b[i] = bias; + } +} + +static void copy_projection( + const MappedProj& src, + const ProjectionShape& shape, + const std::vector& ids, + MTL::Buffer* wbuf, + MTL::Buffer* sbuf, + MTL::Buffer* bbuf) { + auto* wd = static_cast(wbuf->contents()); + auto* sd = static_cast(sbuf->contents()); + auto* bd = static_cast(bbuf->contents()); + auto* ws = static_cast(src.w); + auto* ss = static_cast(src.s); + auto* bs = static_cast(src.b); + for (size_t local = 0; local < ids.size(); ++local) { + size_t expert = static_cast(ids[local]); + std::memcpy(wd + local * shape.weight_per, ws + expert * shape.weight_per, shape.weight_per); + std::memcpy(sd + local * shape.scale_per, ss + expert * shape.scale_per, shape.scale_per); + std::memcpy(bd + local * shape.scale_per, bs + expert * shape.scale_per, shape.scale_per); + } +} + +static std::string metal_source() { + return R"( + #include + using namespace metal; + struct FusedParams { int experts; int hidden; int inter; int group_size; int bits; int k; }; + inline float bf16_to_float(ushort v) { uint u = uint(v) << 16; return as_type(u); } + inline float qvalue(const device uint* w, const device ushort* s, const device ushort* b, + uint expert, uint row, uint col, uint out_dim, uint in_dim, + uint group_size, uint bits) { + uint words_per_row = (in_dim * bits) / 32; + uint groups_per_row = in_dim / group_size; + uint bit_offset = col * bits; + uint word_idx = bit_offset / 32; + uint shift = bit_offset % 32; + uint base = (expert * out_dim + row) * words_per_row; + uint q = w[base + word_idx] >> shift; + if (shift + bits > 32) q |= w[base + word_idx + 1] << (32 - shift); + q &= (1u << bits) - 1u; + uint sb = (expert * out_dim + row) * groups_per_row + col / group_size; + return float(q) * bf16_to_float(s[sb]) + bf16_to_float(b[sb]); + } + kernel void synthetic_moe( + const device float* x [[buffer(0)]], + device float* y [[buffer(1)]], + const device float* scores [[buffer(2)]], + constant FusedParams& p [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + uint token = gid / p.hidden; + uint col = gid % p.hidden; + float acc = 0.0f; + for (int j = 0; j < p.k; ++j) { + acc += scores[token * p.k + j]; + } + y[gid] = x[token * p.hidden + col] * acc; + } + kernel void fused_moe_pairs( + const device float* x [[buffer(0)]], + const device uint* gate_w [[buffer(1)]], const device ushort* gate_s [[buffer(2)]], const device ushort* gate_b [[buffer(3)]], + const device uint* up_w [[buffer(4)]], const device ushort* up_s [[buffer(5)]], const device ushort* up_b [[buffer(6)]], + const device uint* down_w [[buffer(7)]], const device ushort* down_s [[buffer(8)]], const device ushort* down_b [[buffer(9)]], + device float* y [[buffer(10)]], constant FusedParams& p [[buffer(11)]], + uint tid [[thread_position_in_threadgroup]], uint gid [[thread_position_in_grid]]) { + constexpr uint block_size = 512; + constexpr uint lanes_per_row = 16; + constexpr uint rows_per_step = block_size / lanes_per_row; + uint expert = gid / block_size; + if (expert >= uint(p.experts)) return; + uint token = expert / p.k; + uint local_row = tid / lanes_per_row; + uint row_lane = tid % lanes_per_row; + threadgroup float act[1024]; + threadgroup float gate_part[block_size]; + threadgroup float up_part[block_size]; + for (uint row_base = 0; row_base < uint(p.inter); row_base += rows_per_step) { + uint row = row_base + local_row; + float gate_acc = 0.0f, up_acc = 0.0f; + if (row < uint(p.inter)) { + for (uint col = row_lane; col < uint(p.hidden); col += lanes_per_row) { + float xv = x[token * p.hidden + col]; + gate_acc += qvalue(gate_w, gate_s, gate_b, expert, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv; + up_acc += qvalue(up_w, up_s, up_b, expert, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv; + } + } + gate_part[tid] = gate_acc; + up_part[tid] = up_acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (row_lane < stride) { gate_part[tid] += gate_part[tid + stride]; up_part[tid] += up_part[tid + stride]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row < uint(p.inter)) { + float g = gate_part[tid]; + act[row] = g * (1.0f / (1.0f + exp(-g))) * up_part[tid]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint row_base = 0; row_base < uint(p.hidden); row_base += rows_per_step) { + uint row = row_base + local_row; + float acc = 0.0f; + if (row < uint(p.hidden)) { + for (uint col = row_lane; col < uint(p.inter); col += lanes_per_row) { + acc += qvalue(down_w, down_s, down_b, expert, row, col, p.hidden, p.inter, p.group_size, p.bits) * act[col]; + } + } + gate_part[tid] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (row_lane < stride) gate_part[tid] += gate_part[tid + stride]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row < uint(p.hidden)) { + y[expert * p.hidden + row] = gate_part[tid]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + kernel void fused_moe_slots_pairs( + const device float* x [[buffer(0)]], + const device uint* gate_w [[buffer(1)]], const device ushort* gate_s [[buffer(2)]], const device ushort* gate_b [[buffer(3)]], + const device uint* up_w [[buffer(4)]], const device ushort* up_s [[buffer(5)]], const device ushort* up_b [[buffer(6)]], + const device uint* down_w [[buffer(7)]], const device ushort* down_s [[buffer(8)]], const device ushort* down_b [[buffer(9)]], + device float* y [[buffer(10)]], constant FusedParams& p [[buffer(11)]], + const device uint* local_slots [[buffer(12)]], + uint tid [[thread_position_in_threadgroup]], uint gid [[thread_position_in_grid]]) { + constexpr uint block_size = 512; + constexpr uint lanes_per_row = 16; + constexpr uint rows_per_step = block_size / lanes_per_row; + uint active_idx = gid / block_size; + if (active_idx >= uint(p.experts)) return; + uint slot = local_slots[active_idx]; + uint token = active_idx / p.k; + uint local_row = tid / lanes_per_row; + uint row_lane = tid % lanes_per_row; + threadgroup float act[1024]; + threadgroup float gate_part[block_size]; + threadgroup float up_part[block_size]; + for (uint row_base = 0; row_base < uint(p.inter); row_base += rows_per_step) { + uint row = row_base + local_row; + float gate_acc = 0.0f, up_acc = 0.0f; + if (row < uint(p.inter)) { + for (uint col = row_lane; col < uint(p.hidden); col += lanes_per_row) { + float xv = x[token * p.hidden + col]; + gate_acc += qvalue(gate_w, gate_s, gate_b, slot, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv; + up_acc += qvalue(up_w, up_s, up_b, slot, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv; + } + } + gate_part[tid] = gate_acc; + up_part[tid] = up_acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (row_lane < stride) { gate_part[tid] += gate_part[tid + stride]; up_part[tid] += up_part[tid + stride]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row < uint(p.inter)) { + float g = gate_part[tid]; + act[row] = g * (1.0f / (1.0f + exp(-g))) * up_part[tid]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint row_base = 0; row_base < uint(p.hidden); row_base += rows_per_step) { + uint row = row_base + local_row; + float acc = 0.0f; + if (row < uint(p.hidden)) { + for (uint col = row_lane; col < uint(p.inter); col += lanes_per_row) { + acc += qvalue(down_w, down_s, down_b, slot, row, col, p.hidden, p.inter, p.group_size, p.bits) * act[col]; + } + } + gate_part[tid] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) { + if (row_lane < stride) gate_part[tid] += gate_part[tid + stride]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (row_lane == 0 && row < uint(p.hidden)) { + y[active_idx * p.hidden + row] = gate_part[tid]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + kernel void reduce_pairs( + const device float* pair_y [[buffer(0)]], + const device float* scores [[buffer(1)]], + device float* out [[buffer(2)]], + constant FusedParams& p [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + uint token = gid / p.hidden; + uint col = gid % p.hidden; + float acc = 0.0f; + for (int j = 0; j < p.k; ++j) { + uint pair = token * p.k + uint(j); + acc += pair_y[pair * p.hidden + col] * scores[pair]; + } + out[gid] = acc; + } + )"; +} + +class FusedMoePrimitive : public mx::Primitive { + public: + FusedMoePrimitive( + mx::Stream stream, + std::string compute_dir, + int layer, + int hidden, + int inter, + int group, + int bits, + int num_experts, + bool synthetic, + std::vector expert_ids) + : Primitive(stream), + compute_dir_(std::move(compute_dir)), + layer_(layer), + hidden_(hidden), + inter_(inter), + group_(group), + bits_(bits), + num_experts_(num_experts), + synthetic_(synthetic), + expert_ids_(std::move(expert_ids)) { + if (!synthetic_) { + gate_ = map_proj(compute_dir_, layer_, "gate_proj"); + up_ = map_proj(compute_dir_, layer_, "up_proj"); + down_ = map_proj(compute_dir_, layer_, "down_proj"); + } + } + + ~FusedMoePrimitive() override { + if (!synthetic_) { + unmap_proj(gate_); + unmap_proj(up_); + unmap_proj(down_); + } + } + + const char* name() const override { return "FusedMoePrimitive"; } + + void eval_cpu(const std::vector&, std::vector&) override { + throw std::runtime_error("FusedMoePrimitive only supports GPU evaluation"); + } + + void eval_gpu(const std::vector& inputs, std::vector& outputs) override { + auto& x = inputs[0]; + auto& scores = inputs[1]; + auto& out = outputs[0]; + int active = static_cast(expert_ids_.size()); + int tokens = static_cast(x.size() / hidden_); + int k = active / std::max(1, tokens); + if (active <= 0 || tokens <= 0 || active != static_cast(scores.size())) { + throw std::runtime_error("FusedMoePrimitive shape mismatch"); + } + out.set_data(mx::allocator::malloc(out.nbytes())); + auto& d = mx::metal::device(stream().device); + auto& enc = mx::metal::get_command_encoder(stream()); + auto lib = d.get_library("native_moe_mlx_ext", []() { return metal_source(); }); + if (synthetic_) { + auto synthetic_kernel = d.get_kernel("synthetic_moe", lib); + enc.set_compute_pipeline_state(synthetic_kernel); + enc.set_input_array(x, 0); + enc.set_output_array(out, 1); + enc.set_input_array(scores, 2); + FusedParams params{active, hidden_, inter_, group_, bits_, k}; + enc.set_bytes(params, 3); + enc.dispatch_threads(MTL::Size(out.size(), 1, 1), MTL::Size(std::min(out.size(), 256), 1, 1)); + return; + } + auto fused = d.get_kernel("fused_moe_pairs", lib); + auto reduce = d.get_kernel("reduce_pairs", lib); + + ProjectionShape gu = proj_shape(inter_, hidden_, group_, bits_); + ProjectionShape down = proj_shape(hidden_, inter_, group_, bits_); + size_t gu_w_bytes = static_cast(active) * gu.weight_per; + size_t gu_s_bytes = static_cast(active) * gu.scale_per; + size_t down_w_bytes = static_cast(active) * down.weight_per; + size_t down_s_bytes = static_cast(active) * down.scale_per; + buffers_.clear(); + auto make_buffer = [&](size_t nbytes) -> MTL::Buffer* { + buffers_.push_back(NS::TransferPtr(d.mtl_device()->newBuffer(nbytes, MTL::ResourceStorageModeShared))); + return buffers_.back().get(); + }; + MTL::Buffer* gw = make_buffer(gu_w_bytes); + MTL::Buffer* gs = make_buffer(gu_s_bytes); + MTL::Buffer* gb = make_buffer(gu_s_bytes); + MTL::Buffer* uw = make_buffer(gu_w_bytes); + MTL::Buffer* us = make_buffer(gu_s_bytes); + MTL::Buffer* ub = make_buffer(gu_s_bytes); + MTL::Buffer* dw = make_buffer(down_w_bytes); + MTL::Buffer* ds = make_buffer(down_s_bytes); + MTL::Buffer* db = make_buffer(down_s_bytes); + if (synthetic_) { + fill_synthetic(gw, gs, gb, 1); + fill_synthetic(uw, us, ub, 2); + fill_synthetic(dw, ds, db, 3); + } else { + copy_projection(gate_, gu, expert_ids_, gw, gs, gb); + copy_projection(up_, gu, expert_ids_, uw, us, ub); + copy_projection(down_, down, expert_ids_, dw, ds, db); + } + + auto pair_out = mx::zeros(mx::Shape{active, hidden_}, mx::float32, stream()); + pair_out.set_data(mx::allocator::malloc(pair_out.nbytes())); + + enc.set_compute_pipeline_state(fused); + enc.set_input_array(x, 0); + enc.set_buffer(gw, 1); + enc.set_buffer(gs, 2); + enc.set_buffer(gb, 3); + enc.set_buffer(uw, 4); + enc.set_buffer(us, 5); + enc.set_buffer(ub, 6); + enc.set_buffer(dw, 7); + enc.set_buffer(ds, 8); + enc.set_buffer(db, 9); + enc.set_output_array(out, 10); + FusedParams params{active, hidden_, inter_, group_, bits_, k}; + enc.set_bytes(params, 11); + enc.set_input_array(scores, 12); + enc.add_temporary(x); + enc.add_temporary(scores); + enc.dispatch_threads(MTL::Size(static_cast(active) * 512, 1, 1), MTL::Size(512, 1, 1)); + } + + private: + std::string compute_dir_; + int layer_; + int hidden_; + int inter_; + int group_; + int bits_; + int num_experts_; + bool synthetic_; + std::vector expert_ids_; + MappedProj gate_; + MappedProj up_; + MappedProj down_; + std::vector> buffers_; +}; + +class StagedFusedMoePrimitive : public mx::Primitive { + public: + StagedFusedMoePrimitive(mx::Stream stream, int hidden, int inter, int group, int bits) + : Primitive(stream), hidden_(hidden), inter_(inter), group_(group), bits_(bits) {} + + const char* name() const override { return "StagedFusedMoePrimitive"; } + + void eval_cpu(const std::vector&, std::vector&) override { + throw std::runtime_error("StagedFusedMoePrimitive only supports GPU evaluation"); + } + + void eval_gpu(const std::vector& inputs, std::vector& outputs) override { + auto& x = inputs[0]; + auto& scores = inputs[1]; + auto& out = outputs[0]; + int tokens = static_cast(x.size() / hidden_); + int active = static_cast(scores.size()); + int k = active / std::max(1, tokens); + if (active <= 0 || tokens <= 0 || active != static_cast(scores.size())) { + throw std::runtime_error("StagedFusedMoePrimitive shape mismatch"); + } + out.set_data(mx::allocator::malloc(out.nbytes())); + auto& d = mx::metal::device(stream().device); + auto& enc = mx::metal::get_command_encoder(stream()); + auto lib = d.get_library("native_moe_mlx_ext", []() { return metal_source(); }); + auto fused = d.get_kernel("fused_moe_pairs", lib); + auto reduce = d.get_kernel("reduce_pairs", lib); + auto pair_out = mx::zeros(mx::Shape{active, hidden_}, mx::float32, stream()); + pair_out.set_data(mx::allocator::malloc(pair_out.nbytes())); + enc.set_compute_pipeline_state(fused); + enc.set_input_array(x, 0); + enc.set_input_array(inputs[2], 1); + enc.set_input_array(inputs[3], 2); + enc.set_input_array(inputs[4], 3); + enc.set_input_array(inputs[5], 4); + enc.set_input_array(inputs[6], 5); + enc.set_input_array(inputs[7], 6); + enc.set_input_array(inputs[8], 7); + enc.set_input_array(inputs[9], 8); + enc.set_input_array(inputs[10], 9); + enc.set_output_array(pair_out, 10); + FusedParams params{active, hidden_, inter_, group_, bits_, k}; + enc.set_bytes(params, 11); + enc.dispatch_threads(MTL::Size(static_cast(active) * 512, 1, 1), MTL::Size(512, 1, 1)); + enc.set_compute_pipeline_state(reduce); + enc.set_input_array(pair_out, 0); + enc.set_input_array(scores, 1); + enc.set_output_array(out, 2); + enc.set_bytes(params, 3); + enc.add_temporary(pair_out); + enc.dispatch_threads(MTL::Size(out.size(), 1, 1), MTL::Size(std::min(out.size(), 256), 1, 1)); + } + + private: + int hidden_; + int inter_; + int group_; + int bits_; +}; + +class SlotFusedMoePrimitive : public mx::Primitive { + public: + SlotFusedMoePrimitive(mx::Stream stream, int hidden, int inter, int group, int bits) + : Primitive(stream), hidden_(hidden), inter_(inter), group_(group), bits_(bits) {} + + const char* name() const override { return "SlotFusedMoePrimitive"; } + + void eval_cpu(const std::vector&, std::vector&) override { + throw std::runtime_error("SlotFusedMoePrimitive only supports GPU evaluation"); + } + + void eval_gpu(const std::vector& inputs, std::vector& outputs) override { + auto& x = inputs[0]; + auto& local_slots = inputs[1]; + auto& scores = inputs[2]; + auto& out = outputs[0]; + int tokens = static_cast(x.size() / hidden_); + int active = static_cast(scores.size()); + int k = active / std::max(1, tokens); + if (active <= 0 || tokens <= 0 || active != static_cast(local_slots.size())) { + throw std::runtime_error("SlotFusedMoePrimitive shape mismatch"); + } + out.set_data(mx::allocator::malloc(out.nbytes())); + auto& d = mx::metal::device(stream().device); + auto& enc = mx::metal::get_command_encoder(stream()); + auto lib = d.get_library("native_moe_mlx_ext", []() { return metal_source(); }); + auto fused = d.get_kernel("fused_moe_slots_pairs", lib); + auto reduce = d.get_kernel("reduce_pairs", lib); + auto pair_out = mx::zeros(mx::Shape{active, hidden_}, mx::float32, stream()); + pair_out.set_data(mx::allocator::malloc(pair_out.nbytes())); + enc.set_compute_pipeline_state(fused); + enc.set_input_array(x, 0); + enc.set_input_array(inputs[3], 1); + enc.set_input_array(inputs[4], 2); + enc.set_input_array(inputs[5], 3); + enc.set_input_array(inputs[6], 4); + enc.set_input_array(inputs[7], 5); + enc.set_input_array(inputs[8], 6); + enc.set_input_array(inputs[9], 7); + enc.set_input_array(inputs[10], 8); + enc.set_input_array(inputs[11], 9); + enc.set_output_array(pair_out, 10); + FusedParams params{active, hidden_, inter_, group_, bits_, k}; + enc.set_bytes(params, 11); + enc.set_input_array(local_slots, 12); + enc.dispatch_threads(MTL::Size(static_cast(active) * 512, 1, 1), MTL::Size(512, 1, 1)); + enc.set_compute_pipeline_state(reduce); + enc.set_input_array(pair_out, 0); + enc.set_input_array(scores, 1); + enc.set_output_array(out, 2); + enc.set_bytes(params, 3); + enc.add_temporary(pair_out); + enc.dispatch_threads(MTL::Size(out.size(), 1, 1), MTL::Size(std::min(out.size(), 256), 1, 1)); + } + + private: + int hidden_; + int inter_; + int group_; + int bits_; +}; +mx::array fused_moe( + const mx::array& x, + const mx::array& expert_ids, + const mx::array& scores, + const std::string& compute_dir, + int layer, + int hidden, + int inter, + int group, + int bits, + int num_experts, + bool synthetic, + mx::StreamOrDevice s = {}) { + if (!synthetic) { + throw std::runtime_error( + "real mmap staging in MLX Primitive needs MLX-managed staging buffers; use synthetic for bridge tests"); + } + auto ids = expert_ids; + ids.eval(); + std::vector expert_vec(ids.size()); + const uint32_t* idp = ids.data(); + for (size_t i = 0; i < ids.size(); ++i) expert_vec[i] = static_cast(idp[i]); + auto out_shape = x.shape(); + out_shape.back() = hidden; + return mx::array( + out_shape, + mx::float32, + std::make_shared( + mx::to_stream(s), + compute_dir, + layer, + hidden, + inter, + group, + bits, + num_experts, + synthetic, + std::move(expert_vec)), + std::vector{x, scores}); +} + +mx::array fused_moe_staged( + const mx::array& x, + const mx::array& scores, + const mx::array& gate_w, + const mx::array& gate_s, + const mx::array& gate_b, + const mx::array& up_w, + const mx::array& up_s, + const mx::array& up_b, + const mx::array& down_w, + const mx::array& down_s, + const mx::array& down_b, + int hidden, + int inter, + int group, + int bits, + mx::StreamOrDevice s = {}) { + auto out_shape = x.shape(); + out_shape.back() = hidden; + return mx::array( + out_shape, + mx::float32, + std::make_shared(mx::to_stream(s), hidden, inter, group, bits), + std::vector{ + x, scores, + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b}); +} + +mx::array fused_moe_slots( + const mx::array& x, + const mx::array& local_slots, + const mx::array& scores, + const mx::array& gate_w, + const mx::array& gate_s, + const mx::array& gate_b, + const mx::array& up_w, + const mx::array& up_s, + const mx::array& up_b, + const mx::array& down_w, + const mx::array& down_s, + const mx::array& down_b, + int hidden, + int inter, + int group, + int bits, + mx::StreamOrDevice s = {}) { + auto out_shape = x.shape(); + out_shape.back() = hidden; + return mx::array( + out_shape, + mx::float32, + std::make_shared(mx::to_stream(s), hidden, inter, group, bits), + std::vector{ + x, local_slots, scores, + gate_w, gate_s, gate_b, + up_w, up_s, up_b, + down_w, down_s, down_b}); +} diff --git a/native/ext/compute/fused_moe.h b/native/ext/compute/fused_moe.h new file mode 100644 index 0000000..23aac04 --- /dev/null +++ b/native/ext/compute/fused_moe.h @@ -0,0 +1,23 @@ +// fused MoE 计算：合成/真实 mmap staging、staged(预切片权重)、slot(local slot 索引)三条工厂。 +// 实现见 fused_moe.cpp；这里只暴露给 bindings 的自由函数声明（默认参数放在声明处）。 +#pragma once +#include "../common.h" + +mx::array fused_moe( + const mx::array& x, const mx::array& expert_ids, const mx::array& scores, + const std::string& compute_dir, int layer, int hidden, int inter, int group, + int bits, int num_experts, bool synthetic, mx::StreamOrDevice s); + +mx::array fused_moe_staged( + const mx::array& x, const mx::array& scores, + const mx::array& gate_w, const mx::array& gate_s, const mx::array& gate_b, + const mx::array& up_w, const mx::array& up_s, const mx::array& up_b, + const mx::array& down_w, const mx::array& down_s, const mx::array& down_b, + int hidden, int inter, int group, int bits, mx::StreamOrDevice s); + +mx::array fused_moe_slots( + const mx::array& x, const mx::array& local_slots, const mx::array& scores, + const mx::array& gate_w, const mx::array& gate_s, const mx::array& gate_b, + const mx::array& up_w, const mx::array& up_s, const mx::array& up_b, + const mx::array& down_w, const mx::array& down_s, const mx::array& down_b, + int hidden, int inter, int group, int bits, mx::StreamOrDevice s); diff --git a/native/ext/io/bg_reader.cpp b/native/ext/io/bg_reader.cpp new file mode 100644 index 0000000..9969656 --- /dev/null +++ b/native/ext/io/bg_reader.cpp @@ -0,0 +1,183 @@ +// [7] 自由后台读线程（de-risk）：双队列 + 低优并发上限的后台 pread 线程池。 +#include "bg_reader.h" +#include "blob_io.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +struct ReadOp { uint8_t* dst; size_t nbytes; off_t file_off; }; +struct BgJob { + std::vector ops; + std::string path; + long ticket; + std::vector keep; // 持 buffer 引用，保证 dst 指针在读期间存活 + std::function task; // 若设置，worker 直接执行它（侧区异步读用），不走 ops/ticket + int prio = 0; // >0=高优(route 读)，=0=低优(投机兜底) + bool nocache = true; // demand（route）读设 false 走 page cache：段偏移非页对齐， + // F_NOCACHE 下多 MB 非对齐 pread 会短读 → 池槽脏字节（实测）。 +}; + +// 双队列 + 低优并发上限：高优(route)读永远优先取、且独占大多数 worker； +// 低优(投机)读最多占 low_cap_ 个 worker → 给 route 读留出 worker 和 SSD 带宽， +// 实现"高分抢带宽先就绪、低分延后不阻塞关键路径"的真 IO 优先级。 +class BgReader { + public: + ~BgReader() { stop(); } // 进程退出兜底：避免 joinable 线程析构触发 std::terminate + void start(int workers, int low_cap) { + { std::lock_guard lk(m_); low_cap_ = low_cap; } + ensure_started(workers); + } + // 懒启动：侧区预取的 GPU 回调可能在 Python 没调用 bg_reader_start 时就要派任务， + // 故 submit/submit_task 自动确保线程已起（默认 4 worker），调用方无需显式 start。 + void ensure_started(int workers) { + std::lock_guard lk(m_); + if (running_) return; + running_ = true; + for (int i = 0; i < workers; ++i) threads_.emplace_back([this] { loop(); }); + } + void submit(BgJob job) { + ensure_started(4); + { std::lock_guard lk(m_); + if (job.prio > 0) high_q_.push(std::move(job)); else low_q_.push(std::move(job)); } + cv_.notify_all(); + } + void submit_task(std::function fn) { + ensure_started(4); + BgJob job; + job.task = std::move(fn); + { std::lock_guard lk(m_); low_q_.push(std::move(job)); } // 通用任务走低优 + cv_.notify_all(); + } + bool ready(long ticket) { + std::lock_guard lk(dm_); + return done_.count(ticket) > 0; + } + void wait(long ticket) { + std::unique_lock lk(dm_); + dcv_.wait(lk, [&] { return done_.count(ticket) > 0; }); + } + void stop() { + { std::lock_guard lk(m_); running_ = false; } + cv_.notify_all(); + for (auto& t : threads_) if (t.joinable()) t.join(); + threads_.clear(); + { std::lock_guard lk(dm_); done_.clear(); } + { std::lock_guard lk(m_); active_low_ = 0; } + } + private: + int low_budget() const { return low_cap_ <= 0 ? 1 << 30 : low_cap_; } // <=0 视为不限流 + bool can_take_low() const { return !low_q_.empty() && active_low_ < low_budget(); } + void loop() { + std::unordered_map fds; // 每线程各自缓存 fd + while (true) { + std::unique_lock lk(m_); + cv_.wait(lk, [this] { return !high_q_.empty() || can_take_low() || !running_; }); + if (!running_ && high_q_.empty() && low_q_.empty()) break; + if (!running_ && high_q_.empty() && !can_take_low()) break; // 退出时低优可超额排空 + bool is_low = false; + BgJob job; + if (!high_q_.empty()) { // 高优(route)读永远先取 + job = std::move(high_q_.front()); high_q_.pop(); + } else if (can_take_low()) { // 低优限流：仅在额度内取 + job = std::move(low_q_.front()); low_q_.pop(); ++active_low_; is_low = true; + } else { + continue; // 假醒（低优已满额）→ 回去等 + } + lk.unlock(); + if (job.task) { job.task(); } // 通用任务（侧区异步读）：直接执行，无 ticket + else { + int fd; + std::string key = (job.nocache ? "N:" : "C:") + job.path; // 按 nocache 分别缓存 fd + auto it = fds.find(key); + if (it == fds.end()) { + fd = job.nocache ? open_blob_nocache(job.path.c_str()) : ::open(job.path.c_str(), O_RDONLY); + fds[key] = fd; + } else fd = it->second; + if (fd >= 0) + for (auto& op : job.ops) { + ssize_t got = ::pread(fd, op.dst, op.nbytes, op.file_off); + if (got != static_cast(op.nbytes)) + fprintf(stderr, "[bg pread SHORT] got=%zd want=%zu off=%lld\n", + got, op.nbytes, static_cast(op.file_off)); + } + { std::lock_guard lk2(dm_); done_.insert(job.ticket); } + dcv_.notify_all(); + } + if (is_low) { // 释放低优额度 → 唤醒别的 worker 再取 + std::lock_guard lk2(m_); --active_low_; cv_.notify_all(); + } + } + for (auto& kv : fds) if (kv.second >= 0) ::close(kv.second); + } + std::mutex m_, dm_; + std::condition_variable cv_, dcv_; + std::queue high_q_, low_q_; + int active_low_ = 0; // 当前正在执行的低优读数 + int low_cap_ = 0; // 低优并发上限（<=0 不限流，保持旧行为） + std::unordered_set done_; + std::vector threads_; + bool running_ = false; +}; +BgReader g_bg; +} // namespace + +void bg_submit_task(std::function fn) { g_bg.submit_task(std::move(fn)); } + +void bg_reader_start(int workers, int low_cap) { g_bg.start(workers, low_cap); } + +long bg_reader_submit(const mx::array& dst, const std::vector& experts, + const std::vector& rows, const std::string& path, + int stride, long ticket, int prio) { + mx::array d = dst; + d.eval(); + uint8_t* base = d.data(); + size_t st = static_cast(stride); + BgJob job; + job.path = path; + job.ticket = ticket; + job.prio = prio; + job.keep.push_back(d); + for (size_t i = 0; i < experts.size(); ++i) + job.ops.push_back(ReadOp{base + static_cast(rows[i]) * st, st, + static_cast(static_cast(experts[i]) * st)}); + g_bg.submit(std::move(job)); + return ticket; +} + +// 专家各段直写进多个池段张量的 slot 行（消费侧零 MLX 算子）。 +long bg_pread_into_pool( + const std::vector& dst, + const std::vector& seg_off, + const std::vector& seg_nb, + long slot, long expert, + const std::string& path, long stride, long ticket, int prio, bool nocache) { + BgJob job; + job.path = path; + job.ticket = ticket; + job.prio = prio; + job.nocache = nocache; + for (size_t i = 0; i < dst.size(); ++i) { + mx::array d = dst[i]; + d.eval(); + uint8_t* base = d.data(); + job.keep.push_back(d); + job.ops.push_back(ReadOp{ + base + static_cast(slot) * static_cast(seg_nb[i]), + static_cast(seg_nb[i]), + static_cast(static_cast(expert) * static_cast(stride) + + static_cast(seg_off[i]))}); + } + g_bg.submit(std::move(job)); + return ticket; +} + +bool bg_reader_ready(long ticket) { return g_bg.ready(ticket); } +void bg_reader_wait(long ticket) { g_bg.wait(ticket); } +void bg_reader_stop() { g_bg.stop(); } diff --git a/native/ext/io/bg_reader.h b/native/ext/io/bg_reader.h new file mode 100644 index 0000000..779ce87 --- /dev/null +++ b/native/ext/io/bg_reader.h @@ -0,0 +1,23 @@ +// 自由后台读线程（de-risk）：脱离 GPU 完成回调、零 GIL，pread 进调用方 MLX buffer。 +// 线程全程只碰 dst 原始指针 / 路径 / 整数，绝不接触 Python 对象 → 无需 GIL。 +#pragma once +#include "../common.h" +#include + +void bg_reader_start(int workers, int low_cap = 0); +long bg_reader_submit(const mx::array& dst, const std::vector& experts, + const std::vector& rows, const std::string& path, + int stride, long ticket, int prio = 0); +bool bg_reader_ready(long ticket); +void bg_reader_wait(long ticket); +void bg_reader_stop(); +long bg_pread_into_pool( + const std::vector& dst, + const std::vector& seg_off, + const std::vector& seg_nb, + long slot, long expert, + const std::string& path, long stride, long ticket, int prio = 0, bool nocache = true); + +// 通用后台任务入口：把任意闭包派到后台线程池（低优队列）。侧区/staging 预取的 GPU 完成回调 +// 用它把 pread 派离 Metal 回调线程 → 真正与计算重叠。内部接口，不绑 Python。 +void bg_submit_task(std::function fn); diff --git a/native/ext/io/blob_io.h b/native/ext/io/blob_io.h new file mode 100644 index 0000000..793b496 --- /dev/null +++ b/native/ext/io/blob_io.h @@ -0,0 +1,16 @@ +// 共用底层 IO 助手：打开 blob 文件。多个模块（侧区/后台读线程等）的 pread 都用它。 +#pragma once +#include +#include + +#ifndef F_NOCACHE +#define F_NOCACHE 48 // macOS：提示内核读过的页不留 page cache +#endif + +// 打开 blob 并设 F_NOCACHE：与 demand 侧 blob_loader 对齐，避免预取读把 page cache 灌满 → +// 在内存受限机上累积压力触发"双稳态"慢挡翻转（实测 zerocopy 每轮翻慢挡的根因）。 +static inline int open_blob_nocache(const char* path) { + int fd = ::open(path, O_RDONLY); + if (fd >= 0) ::fcntl(fd, F_NOCACHE, 1); + return fd; +} diff --git a/native/ext/io/blob_load.cpp b/native/ext/io/blob_load.cpp new file mode 100644 index 0000000..efbd389 --- /dev/null +++ b/native/ext/io/blob_load.cpp @@ -0,0 +1,58 @@ +// [1] blob 直读：pread 专家字节进 MLX buffer（惰性图节点）。 +// 把一组专家的 blob 字节直接 pread 进 MLX 自有 buffer（无 kernel、无额外拷贝）。 +// load 作为惰性图节点：在批量 eval 中执行，避免 Python 侧 per-expert mx.eval 同步。 +#include "blob_load.h" + +#include +#include + +class BlobLoadPrimitive : public mx::Primitive { + public: + BlobLoadPrimitive(mx::Stream stream, std::string path, size_t stride, std::vector experts) + : Primitive(stream), path_(std::move(path)), stride_(stride), experts_(std::move(experts)) {} + const char* name() const override { return "BlobLoadPrimitive"; } + void eval_cpu(const std::vector&, std::vector& outputs) override { + load(outputs[0]); + } + void eval_gpu(const std::vector&, std::vector& outputs) override { + load(outputs[0]); + } + + private: + void load(mx::array& out) { + out.set_data(mx::allocator::malloc(out.nbytes())); + uint8_t* dst = out.data(); + int fd = ::open(path_.c_str(), O_RDONLY); + if (fd < 0) throw std::runtime_error("blob open failed: " + path_); + for (size_t i = 0; i < experts_.size(); ++i) { + size_t off = static_cast(experts_[i]) * stride_; + ssize_t n = ::pread(fd, dst + i * stride_, stride_, static_cast(off)); + if (n != static_cast(stride_)) { + ::close(fd); + throw std::runtime_error("blob pread short read"); + } + } + ::close(fd); + } + std::string path_; + size_t stride_; + std::vector experts_; +}; + +mx::array blob_load( + const std::string& path, + const mx::array& expert_ids, + int stride, + mx::StreamOrDevice s = {}) { + auto ids = expert_ids; + ids.eval(); + std::vector ev(ids.size()); + const uint32_t* p = ids.data(); + for (size_t i = 0; i < ids.size(); ++i) ev[i] = static_cast(p[i]); + int n = static_cast(ev.size()); + return mx::array( + mx::Shape{n, stride}, + mx::uint8, + std::make_shared(mx::to_stream(s), path, static_cast(stride), std::move(ev)), + std::vector{}); +} diff --git a/native/ext/io/blob_load.h b/native/ext/io/blob_load.h new file mode 100644 index 0000000..808dfed --- /dev/null +++ b/native/ext/io/blob_load.h @@ -0,0 +1,7 @@ +// blob 直读：把一组专家 blob 字节 pread 进新建的 MLX uint8[n,stride] 数组（惰性图节点）。 +#pragma once +#include "../common.h" + +mx::array blob_load( + const std::string& path, const mx::array& expert_ids, int stride, + mx::StreamOrDevice s); diff --git a/native/ext/pool/demand.cpp b/native/ext/pool/demand.cpp new file mode 100644 index 0000000..bbe53c1 --- /dev/null +++ b/native/ext/pool/demand.cpp @@ -0,0 +1,331 @@ +// [6] Phase 2 方案B：真实区槽状态 C++ 全接管（1 次同步版）。 +// 复刻 Python ResidentExpertPool 的 _slot_of/_free/_freq + _choose_victim/_alloc_slot 语义： +// - free 优先(front pop，free 初始 [0,cap))；free 空 → LFU 驱逐，受害者槽直接复用(不回 free)。 +// - _choose_victim：candidates=插入序中 ∉current 者；victim=min(freq, 候选下标)，驱逐不删 freq。 +// - dual 语义：真实区命中不移动插入序；仅新放入 miss 追加序尾。 +#include "demand.h" +#include "side_region.h" +#include "../io/bg_reader.h" + +#include +#include +#include +#include +#include + +struct RealLayer { + std::vector order; // 插入序(LRU tie-break)，与 e2r 同步维护 + std::unordered_map e2r; // expert -> slot [0,cap) + std::vector free_rows; // 空闲槽(front pop，仿 free.pop(0)) + std::unordered_map freq; // LFU 频次(驱逐不删，与 Python 一致) + std::unordered_set pinned; // AUTOPIN 钉死集：驱逐永不选(见 real_pin) + int cap = 0; + long access = 0; // 累计访问(decay 用) + bool inited = false; +}; +static std::mutex g_real_mutex; +static std::map g_real; + +// demand 统计：累计 + 本次(供 Python 更新 rp.hits/misses/gpu_fastpath/gpu_fallback)。 +static std::mutex g_dstat_mutex; +// 本次 [hitpos, misspos, loads, fallback01, unplaced] +static long g_d_last[5] = {0, 0, 0, 0, 0}; +static std::atomic g_demand_ticket{1000000000}; // demand 并行 pread 用的独立 ticket 段 + +// 诊断计时(DEMAND_TIMING=1)：累计各段主线程微秒，供定位结构性开销。默认关。 +static bool g_dt_on = false; +static double g_dt[6] = {0, 0, 0, 0, 0, 0}; // [inds_eval, pool_eval, side_snap, real_lock, core, build] +static inline double dt_now_us() { + return std::chrono::duration( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} +std::vector demand_timings() { return {g_dt[0], g_dt[1], g_dt[2], g_dt[3], g_dt[4], g_dt[5]}; } +void demand_timing_enable(bool on) { + g_dt_on = on; + for (int i = 0; i < 6; ++i) g_dt[i] = 0; +} + +static void real_ensure_locked(RealLayer& c, int cap) { + if (c.inited) return; + c.cap = cap; + c.free_rows.clear(); + for (int r = 0; r < cap; ++r) c.free_rows.push_back(r); // free 初始 [0,cap) + c.inited = true; +} + +void real_init(int layer, int cap) { + std::lock_guard lk(g_real_mutex); + real_ensure_locked(g_real[layer], cap); +} + +std::vector real_region_contents(int layer) { + std::lock_guard lk(g_real_mutex); + std::vector out; + auto it = g_real.find(layer); + if (it != g_real.end()) + for (auto& p : it->second.e2r) { out.push_back(p.first); out.push_back(p.second); } + return out; +} + +int real_region_count(int layer) { + std::lock_guard lk(g_real_mutex); + auto it = g_real.find(layer); + return it == g_real.end() ? 0 : static_cast(it->second.e2r.size()); +} + +int real_pinned_count(int layer) { + std::lock_guard lk(g_real_mutex); + auto it = g_real.find(layer); + return it == g_real.end() ? 0 : static_cast(it->second.pinned.size()); +} + +void real_reset() { + std::lock_guard lk(g_real_mutex); + g_real.clear(); +} + +// 复刻 _choose_victim：遍历插入序(order)选 ∉current 且 freq 最小者，并列取最早(候选下标最小)。 +// pinned(AUTOPIN 钉死)与 current 同样永不选。返回 expert id；-1 表示无可驱逐。调用方须持 g_real_mutex。 +static int choose_victim_locked(RealLayer& c, const std::unordered_set& current) { + int victim = -1; + uint32_t best = 0; + for (int e : c.order) { + if (!c.e2r.count(e) || current.count(e) || c.pinned.count(e)) continue; + uint32_t f = c.freq.count(e) ? c.freq[e] : 0; + if (victim < 0 || f < best) { victim = e; best = f; } // 并列不更新 → 保留更早者 + } + return victim; +} + +// AUTOPIN 预热钉死：把历史热专家注册进真实区并标记 pinned(此后驱逐永不选它)。 +// 已驻 → 仅补 pinned 标记、返回原槽；未驻 → free 头取槽登记(池行字节由调用方负责写入)。 +// 返回与 experts 平行的槽位；空闲耗尽分不到槽给 -1(调用方跳过该专家)。幂等。 +// 前置：调用方须先 real_init(否则 free_rows 为空、全返 -1)。 +std::vector real_pin(int layer, const std::vector& experts) { + std::lock_guard lk(g_real_mutex); + RealLayer& c = g_real[layer]; + std::vector out; + out.reserve(experts.size()); + for (int e : experts) { + auto it = c.e2r.find(e); + if (it != c.e2r.end()) { c.pinned.insert(e); out.push_back(it->second); continue; } + if (c.free_rows.empty()) { out.push_back(-1); continue; } + int slot = c.free_rows.front(); + c.free_rows.erase(c.free_rows.begin()); // free.pop(0) + c.e2r[e] = slot; + c.order.push_back(e); + c.pinned.insert(e); + out.push_back(slot); + } + return out; +} + +// AUTOPIN 热度持久化：导出各层 LFU 累计频次，扁平 [layer, expert, count, ...]。 +// 仅 lfu 策略下 demand 路径计频；驱逐不删 freq，故即该进程全生命周期的路由热度， +// 由 Python 侧按增量差分合并进 usage 文件(dual decode 无 Python 物化点，这是唯一计频来源)。 +std::vector real_freq_dump() { + std::lock_guard lk(g_real_mutex); + std::vector out; + for (auto& kv : g_real) + for (auto& p : kv.second.freq) { + out.push_back(static_cast(kv.first)); + out.push_back(static_cast(p.first)); + out.push_back(static_cast(p.second)); + } + return out; +} + +// 复刻 _alloc_slot（e 为 miss，不会已在 e2r）：free 优先，否则 LFU 驱逐复用受害者槽。 +// 返回 slot；-1 表示无可驱逐(超容量，不该在 dual 发生)。调用方须持 g_real_mutex。 +static int alloc_slot_locked(RealLayer& c, int e, const std::unordered_set& current) { + auto it = c.e2r.find(e); + if (it != c.e2r.end()) return it->second; + int slot; + if (!c.free_rows.empty()) { + slot = c.free_rows.front(); + c.free_rows.erase(c.free_rows.begin()); // free.pop(0) + } else { + int victim = choose_victim_locked(c, current); + if (victim < 0) return -1; + slot = c.e2r[victim]; + c.e2r.erase(victim); + for (auto oit = c.order.begin(); oit != c.order.end(); ++oit) + if (*oit == victim) { c.order.erase(oit); break; } + } + c.e2r[e] = slot; + c.order.push_back(e); + return slot; +} + +// LFU 频次 bump（canonical：本次全部唯一专家各 +1）+ decay。调用方须持 g_real_mutex。 +static void note_access_locked(RealLayer& c, const std::vector& uniq_access, + bool lfu, int decay_interval) { + if (!lfu) return; + for (int e : uniq_access) c.freq[e] += 1; + c.access += static_cast(uniq_access.size()); + if (decay_interval > 0 && c.access >= decay_interval) { + for (auto it = c.freq.begin(); it != c.freq.end();) { + it->second /= 2; + if (it->second == 0) it = c.freq.erase(it); + else ++it; + } + c.access = 0; + } +} + +// 核心状态机（纯 CPU、无 I/O）：给定 host inds(ip[n]) + 侧区快照(side)，算 local、分配 miss 槽， +// 把需要落盘的新放入 (expert, slot) 追加到 placements（字节由调用方在锁外并行 pread 落池）。 +// 返回 local(int32 vector)。stats: [hitpos, misspos, loads(=placements 数), unplaced]。 +static std::vector demand_core_locked( + RealLayer& c, const uint32_t* ip, size_t n, const std::unordered_map& side, + bool lfu, int decay_interval, std::vector>& placements, long stats[4]) { + std::vector local(n, -1); + std::vector uniq_miss, access_order; + std::unordered_set miss_seen, access_seen; + int hitpos = 0; + // pass1：算命中(侧区覆盖真实区)、收集 miss(首见序)、收集唯一访问(freq)。 + for (size_t i = 0; i < n; ++i) { + int e = static_cast(ip[i]); + if (access_seen.insert(e).second) access_order.push_back(e); + auto sit = side.find(e); + if (sit != side.end()) { local[i] = sit->second; ++hitpos; continue; } + auto rit = c.e2r.find(e); + if (rit != c.e2r.end()) { local[i] = rit->second; ++hitpos; continue; } + if (miss_seen.insert(e).second) uniq_miss.push_back(e); + } + note_access_locked(c, access_order, lfu, decay_interval); + // pass2：miss 分配槽（不落盘）。current = 本前向全部唯一路由专家(命中+miss)：绝不驱逐本前向要读的 + // 任何专家的槽（否则真实区命中专家的槽被 miss 复写 → 脏字节。比 Python 仅护 miss 更严格、更正确）。 + std::unordered_map new_slot; + long unplaced = 0; + for (int e : uniq_miss) { + int slot = alloc_slot_locked(c, e, access_seen); + if (slot < 0) { + // 超容量：可驱逐槽(= cap − pinned − 本前向命中数)不够安放全部 miss。只能落 0 号槽， + // 即拿别的专家的权重参与计算 → 该前向逐位错算。调用方须靠 stats[3] 拦截并回退到 + // host/fetch 路径；这里只保证「绝不静默」。 + new_slot[e] = 0; + ++unplaced; + continue; + } + new_slot[e] = slot; + placements.emplace_back(e, slot); + } + // pass3：回填 miss 位置。 + int misspos = 0; + for (size_t i = 0; i < n; ++i) { + if (local[i] < 0) { + auto it = new_slot.find(static_cast(ip[i])); + local[i] = (it != new_slot.end()) ? it->second : 0; + ++misspos; + } + } + stats[0] = hitpos; stats[1] = misspos; stats[2] = static_cast(placements.size()); + stats[3] = unplaced; + return local; +} + +// demand 全接管：inds 惰性(内部 eval 一次=1 次同步)；side_gen 指定侧区代；pool_list 为 _segs 顺序 +// 的 per-key 池数组(已 eval、指针稳定)。返回 local(int32, inds.shape)。 +mx::array demand_dual( + const mx::array& inds, const std::vector& pool_list, + const std::vector& seg_nbytes, int layer, int side_gen, const std::string& path, + int stride, int cap, bool lfu, int decay_interval, mx::StreamOrDevice s = {}) { + if (seg_nbytes.size() != pool_list.size()) + throw std::invalid_argument("demand_dual: seg_nbytes.size() != pool_list.size()"); + double t0 = g_dt_on ? dt_now_us() : 0; + // 关键:inds 常是 argpartition(...)[..., -k:] 切片 → 非连续 strided 视图(每行按父数组 stride E + // 偏移)。若直接按连续读 data(),仅首行(token0)正确,后续 token 读到错位内存 → 装错专家、seq≥2 全错 + // (投机 verify 全体受害)。contiguous() 强制物化为连续,读取才逐位正确。 + mx::array ids = mx::contiguous(inds); + ids.eval(); // 唯一同步 + size_t n = ids.size(); + const uint32_t* ip = ids.data(); + double t1 = g_dt_on ? dt_now_us() : 0; + std::vector ptrs; + ptrs.reserve(pool_list.size()); + for (auto& a0 : pool_list) { mx::array a = a0; a.eval(); ptrs.push_back(a.data()); } + double t2 = g_dt_on ? dt_now_us() : 0; + std::unordered_map side = sideregion_snapshot(layer, side_gen); // 侧区快照(该代) + double ta = g_dt_on ? dt_now_us() : 0; + long stats[4]; + std::vector local; + std::vector> placements; // (expert, slot)：锁外并行落盘 + double tb; + { + std::lock_guard lk(g_real_mutex); + tb = g_dt_on ? dt_now_us() : 0; + RealLayer& c = g_real[layer]; + real_ensure_locked(c, cap); + local = demand_core_locked(c, ip, n, side, lfu, decay_interval, placements, stats); + } + // 锁外：把 miss 专家的字节 pread 落真实区槽。复刻基线并发模型——多 miss 派给 BgReader worker + // 并行 pread（高优队列、直写池段行，无主线程 tmp/memcpy），本层等待其完成（并行 → 远快于串行）。 + static const bool kSkipIO = []() { + const char* e = std::getenv("DEMAND_SKIP_IO"); + return e && e[0] == '1'; + }(); + if (!placements.empty() && !kSkipIO) { + std::vector seg_off(seg_nbytes.size()), seg_nb(seg_nbytes.size()); + long acc = 0; + for (size_t k = 0; k < seg_nbytes.size(); ++k) { + seg_off[k] = acc; seg_nb[k] = seg_nbytes[k]; acc += seg_nbytes[k]; + } + std::vector tickets; + tickets.reserve(placements.size()); + for (auto& pr : placements) { + long tk = g_demand_ticket.fetch_add(1); + bg_pread_into_pool(pool_list, seg_off, seg_nb, pr.second, pr.first, path, + static_cast(stride), tk, /*prio=*/1, + /*nocache=*/false); // demand route 读走 cache：段偏移非页对齐 + tickets.push_back(tk); + } + for (long tk : tickets) bg_reader_wait(tk); // 并行 pread 完成 → 池槽字节就绪 + } + double t3 = g_dt_on ? dt_now_us() : 0; + { + std::lock_guard lk(g_dstat_mutex); + g_d_last[0] = stats[0]; g_d_last[1] = stats[1]; g_d_last[2] = stats[2]; + g_d_last[3] = (stats[1] == 0) ? 0 : 1; + g_d_last[4] = stats[3]; + } + if (stats[3] > 0) { + static std::atomic warned{0}; + long k = warned.fetch_add(1); + if (k < 8 || k % 1000 == 0) + fprintf(stderr, + "[DEMAND_DUAL] 错算：layer=%d 有 %ld 个专家分不到槽(落 0 号槽)。" + "inds=%zu cap=%d pinned=%d —— 该前向输出不可信,请调小 PREFILL_CHUNK " + "或调大 EXPERT_SLOTS / 调小 AUTOPIN_BUDGET_FRAC。\n", + layer, stats[3], n, cap, real_pinned_count(layer)); + } + mx::array out = mx::array(local.data(), ids.shape(), mx::int32); + if (g_dt_on) { + double t4 = dt_now_us(); + g_dt[0] += t1 - t0; g_dt[1] += t2 - t1; g_dt[2] += ta - t2; + g_dt[3] += tb - ta; g_dt[4] += t3 - tb; g_dt[5] += t4 - t3; + } + return out; +} + +// 取本次 demand 统计 [hitpos, misspos, loads, fallback01, unplaced]（主线程串行，安全）。 +std::vector demand_last_stats() { + std::lock_guard lk(g_dstat_mutex); + return {g_d_last[0], g_d_last[1], g_d_last[2], g_d_last[3], g_d_last[4]}; +} + +// 测试壳：纯状态推进(不 pread/不侧区)，把 experts_flat 当一次 demand 的 inds，返回 local 槽位。 +// 供 LFU 驱逐语义逐步等价单测(与 Python 参考对拍)。 +std::vector real_debug_place(int layer, const std::vector& experts_flat, int cap, + bool lfu, int decay_interval) { + std::vector u(experts_flat.begin(), experts_flat.end()); + long stats[4]; + std::vector> placements; + std::lock_guard lk(g_real_mutex); + RealLayer& c = g_real[layer]; + real_ensure_locked(c, cap); + std::unordered_map empty_side; + auto local = demand_core_locked(c, u.data(), u.size(), empty_side, lfu, decay_interval, + placements, stats); + return std::vector(local.begin(), local.end()); +} diff --git a/native/ext/pool/demand.h b/native/ext/pool/demand.h new file mode 100644 index 0000000..9794480 --- /dev/null +++ b/native/ext/pool/demand.h @@ -0,0 +1,29 @@ +// Phase 2 方案B：真实区槽状态 C++ 全接管（1 次同步版）+ demand_dual。 +// 复刻 Python ResidentExpertPool 的 free/LFU 语义，成为 dual-source decode 真实区唯一权威。 +#pragma once +#include "../common.h" + +void real_init(int layer, int cap); // 初始化某层真实区(cap 槽全空闲)，幂等 +std::vector real_region_contents(int layer); // [expert, slot, ...] +int real_region_count(int layer); +void real_reset(); +// AUTOPIN：注册 pinned 热专家(驱逐免选)，返回平行槽位(-1=分不到)；前置 real_init。 +std::vector real_pin(int layer, const std::vector& experts); +// 该层 pinned(永不可驱逐)专家数。调用方据此算「可安放的唯一专家上限」= cap − pinned， +// 用于在超容量前把该前向分流到 host/fetch 路径（超了会落 0 号槽 → 逐位错算）。 +int real_pinned_count(int layer); +// AUTOPIN：导出各层 LFU 累计频次，扁平 [layer, expert, count, ...]。 +std::vector real_freq_dump(); +// demand 全接管：inds 惰性(内部 eval 一次)；side_gen 侧区代；pool_list 为 _segs 顺序池数组。 +mx::array demand_dual( + const mx::array& inds, const std::vector& pool_list, + const std::vector& seg_nbytes, int layer, int side_gen, const std::string& path, + int stride, int cap, bool lfu, int decay_interval, mx::StreamOrDevice s); +// [hitpos, misspos, loads, fallback01, unplaced] +// unplaced：本次分不到槽、被迫落 0 号槽的唯一专家数。>0 即该前向逐位错算(必须为 0)。 +std::vector demand_last_stats(); +std::vector demand_timings(); // [inds_eval, pool_eval, state, build] us +void demand_timing_enable(bool on); +// 测试壳：纯状态推进(不 pread/不侧区)，返回 local 槽位；供 LFU 驱逐等价对拍。 +std::vector real_debug_place(int layer, const std::vector& experts_flat, int cap, + bool lfu, int decay_interval); diff --git a/native/ext/pool/owned_pool.cpp b/native/ext/pool/owned_pool.cpp new file mode 100644 index 0000000..ba0ff1d --- /dev/null +++ b/native/ext/pool/owned_pool.cpp @@ -0,0 +1,89 @@ +// [5] Route 3 Phase 1 底座：C++ 拥有的池 buffer。 +// 用 mx::allocator::malloc 分配 buffer、no-op deleter 建叶子 mx.array（C++ 经 g_owned_bufs +// 持有 Buffer 句柄，进程内不释放）。因 C++ 独占持有、MLX 只读，该 buffer 永不被 MLX +// donation/迁移（spike 已证），故侧区/demand 的后台异步 pread 可安全直写、消费侧读同一块。 +#include "owned_pool.h" + +static std::mutex g_owned_mutex; +static std::vector g_owned_bufs; + +static mx::Dtype dtype_from_str(const std::string& s) { + if (s == "uint32") return mx::uint32; + if (s == "uint16") return mx::uint16; + if (s == "uint8") return mx::uint8; + if (s == "int32") return mx::int32; + if (s == "int16") return mx::int16; + if (s == "bfloat16") return mx::bfloat16; + if (s == "float16") return mx::float16; + if (s == "float32") return mx::float32; + throw std::runtime_error("pool_owned_zeros: unsupported dtype " + s); +} + +mx::array pool_owned_zeros(const std::vector& shape, const std::string& dtype) { + mx::Dtype dt = dtype_from_str(dtype); + mx::Shape shp(shape.begin(), shape.end()); + size_t n = 1; + for (int d : shape) n *= static_cast(d); + size_t nbytes = n * static_cast(mx::size_of(dt)); + auto buf = mx::allocator::malloc(nbytes); + std::memset(buf.raw_ptr(), 0, nbytes); + { + std::lock_guard lk(g_owned_mutex); + g_owned_bufs.push_back(buf); // C++ 持有，保证进程内不被释放 + } + return mx::array(buf, shp, dt, [](mx::allocator::Buffer) {}); +} + +// demand 真实区落池：把一批已加载专家段 memcpy 进 owned 池行（无 MLX scatter → pool buffer 永不重绑）。 +// pool_list[i]：第 i 个 key 的池数组；srcs_flat 长 K*m，key-major：[k0行0,k0行1,...,k1行0,...]； +// slots[j]：第 j 个专家的目标物理行。CPU 端直写，调用点须保证此刻池未被 GPU 并发读该行。 +void pool_write_rows(const std::vector& pool_list, + const std::vector& srcs_flat, + const std::vector& slots) { + int K = static_cast(pool_list.size()); + int m = static_cast(slots.size()); + if (m == 0) return; + if (static_cast(srcs_flat.size()) != K * m) + throw std::runtime_error("pool_write_rows: srcs_flat 长度须 == K*m"); + for (int i = 0; i < K; ++i) { + mx::array p = pool_list[i]; + p.eval(); + uint8_t* base = p.data(); + for (int j = 0; j < m; ++j) { + mx::array s = srcs_flat[static_cast(i) * m + j]; + s.eval(); + size_t nb = s.nbytes(); + std::memcpy(base + static_cast(slots[j]) * nb, s.data(), nb); + } + } +} + +// 同上，但源是每 key 预堆叠的 (m,*shape) 整块 stacked_list[i]，按行 memcpy 到 slots[j]。 +void pool_write_stacked(const std::vector& pool_list, + const std::vector& stacked_list, + const std::vector& slots) { + int K = static_cast(pool_list.size()); + int m = static_cast(slots.size()); + if (m == 0) return; + for (int i = 0; i < K; ++i) { + mx::array p = pool_list[i]; + p.eval(); + mx::array st = stacked_list[i]; + st.eval(); + uint8_t* base = p.data(); + const uint8_t* src = st.data(); + size_t rownb = st.nbytes() / static_cast(m); // stacked 为 (m,*shape) → 每行字节 + for (int j = 0; j < m; ++j) + std::memcpy(base + static_cast(slots[j]) * rownb, + src + static_cast(j) * rownb, rownb); + } +} + +// 临时诊断：返回某 mx.array 底层 buffer 的原始数据指针（uintptr）。用于对拍 +// 「C++ 侧区 memcpy 写入的 buffer 指针」与「Python consume/verify 读到的 pool buffer 指针」 +// 是否同一块——若不同即证明 MLX 在两者之间重分配了池 buffer，raw 写入被落单。 +uintptr_t array_data_ptr(const mx::array& a) { + mx::array b = a; + b.eval(); + return reinterpret_cast(b.data()); +} diff --git a/native/ext/pool/owned_pool.h b/native/ext/pool/owned_pool.h new file mode 100644 index 0000000..d1cb562 --- /dev/null +++ b/native/ext/pool/owned_pool.h @@ -0,0 +1,14 @@ +// Route 3 Phase 1 底座：C++ 拥有的池 buffer + 直写（替代 mx.zeros 建池 + 消费侧 MLX scatter）。 +#pragma once +#include "../common.h" + +// C++ 拥有的池 buffer 包成 mx.array（no-op deleter，进程内持有、永不迁移）。 +// 供侧区/demand 后台 pread 安全直写；替代原 mx.zeros 建池 + 消费侧 MLX scatter。 +mx::array pool_owned_zeros(const std::vector& shape, const std::string& dtype); +// demand 真实区落池：把已加载专家段直接 memcpy 进 owned 池行（无 MLX scatter，保 buffer 稳定）。 +void pool_write_rows(const std::vector& pool_list, + const std::vector& srcs_flat, const std::vector& slots); +void pool_write_stacked(const std::vector& pool_list, + const std::vector& stacked_list, const std::vector& slots); +// 诊断：mx.array 底层 buffer 原始指针（uintptr），用于对拍侧区写入 buffer 与 consume 读到 buffer 是否同一块。 +uintptr_t array_data_ptr(const mx::array& a); diff --git a/native/ext/pool/side_region.cpp b/native/ext/pool/side_region.cpp new file mode 100644 index 0000000..a49337d --- /dev/null +++ b/native/ext/pool/side_region.cpp @@ -0,0 +1,389 @@ +// [4] 段散写持久侧区缓存（zero-copy dual-source 默认路径）。 +// 与持久 staging 缓存同构，但目标是“多个结构化 per-key 池数组”：每个池数组形如 +// (cap+spec, ...)，行 [base_row, base_row+spec) 为侧区。命中缺口时 pread blob 整行， +// 再把该行内按固定顺序拼接的各段 memcpy 进对应 per-key 数组的同一物理行。 +#include "side_region.h" +#include "../io/blob_io.h" +#include "../io/bg_reader.h" + +#include +#include +#include +#include +#include +#include + +struct SideLayer { + std::map e2r; // expert -> 物理侧区行 [base_row, base_row+spec) + std::vector free_rows; + std::map freq; // expert -> 预测频次(LFU 分数;仅 SIDEREGION_LFU 用) + bool inited = false; + int base = 0; // 侧区起始物理行 base_row + int spec = 0; // 侧区行数 spec_slots +}; +static std::mutex g_side_mutex; +static std::map, SideLayer> g_side; // 键 (layer, gen)：双缓冲两代独立 + +// 侧区 fill 在途计数：eval_gpu 提交预取时 +1，后台 read_publish 写完字节后 -1。 +// 消费方在前向开头 sideregion_drain() 排空上一前向的 fill，保证被消费的侧区行字节 +// 已完全写好（含 GPU 完成回调滞后的情形）→ 消灭「GPU 消费 kernel 读到半写侧区行」竞态。 +static std::atomic g_side_inflight{0}; +static std::mutex g_side_drain_mutex; +static std::condition_variable g_side_drain_cv; +static inline void side_inflight_done() { + { + std::lock_guard lk(g_side_drain_mutex); + g_side_inflight.fetch_sub(1); + } + g_side_drain_cv.notify_all(); +} +void sideregion_drain() { + std::unique_lock lk(g_side_drain_mutex); + g_side_drain_cv.wait(lk, [] { return g_side_inflight.load() == 0; }); +} + +// 临时诊断：只追踪指定 (layer,row) 的所有账本/字节事件（SIDE_TRACE_LAYER/ROW）。 +static inline bool side_trace_hit(int layer, int row) { + const char* le = std::getenv("SIDE_TRACE_LAYER"); + const char* re = std::getenv("SIDE_TRACE_ROW"); + if (!le || !re) return false; + return layer == atoi(le) && row == atoi(re); +} + +// 临时诊断：跨线程全局事件序号 + 线程短 id，用于把 reserve(Metal 回调线程)/read_publish(bg 线程)/ +// sideregion_kv(主线程) 的事件按发生顺序排出来（root-cause 取证，SIDE_TRACE_* 命中时才用）。 +static std::atomic g_side_ev{0}; +static inline unsigned side_tid() { + return static_cast( + std::hash{}(std::this_thread::get_id()) & 0xffff); +} + +class PrefetchPoolSideRegionPrimitive : public mx::Primitive { + public: + PrefetchPoolSideRegionPrimitive(mx::Stream s, std::vector seg_nbytes, int layer, int gen, + std::string path, size_t stride, std::vector resident, + int spec_slots, int base_row) + : Primitive(s), seg_(std::move(seg_nbytes)), layer_(layer), gen_(gen), path_(std::move(path)), + stride_(stride), resident_(std::move(resident)), spec_(spec_slots), base_(base_row) {} + const char* name() const override { return "PrefetchPoolSideRegionPrimitive"; } + + void eval_cpu(const std::vector& in, std::vector& out) override { + out[0].set_data(mx::allocator::malloc(out[0].nbytes())); + run(in); + } + void eval_gpu(const std::vector& in, std::vector& out) override { + out[0].set_data(mx::allocator::malloc(out[0].nbytes())); + std::vector ptrs; + ptrs.reserve(in.size() - 1); + for (size_t i = 1; i < in.size(); ++i) { + mx::array a = in[i]; // 非 const 拷贝才能取可写指针 + ptrs.push_back(a.data()); + } + mx::array ids = in[0]; + const uint32_t* idp = ids.data(); + size_t n = ids.size(); + std::vector seg = seg_; + int layer = layer_; + int gen = gen_; + std::string path = path_; + size_t stride = stride_; + std::vector resident = resident_; + int spec = spec_, base = base_; + auto& enc = mx::metal::get_command_encoder(stream()); + MTL::CommandBuffer* cb = enc.get_command_buffer(); + // 提交即计在途：在 eval_gpu（预取提交）时 +1，直到后台字节写完才 -1。这样消费方前向开头 + // sideregion_drain() 能等到「即使 GPU 完成回调尚未触发」的 fill，闭合跨前向的写-读竞态。 + g_side_inflight.fetch_add(1); + // in 按值捕获 → 保活 expert_ids 与所有池数组 buffer；idp/ptrs 在回调里指针有效。 + cb->addCompletedHandler( + [in, ptrs, seg, idp, n, layer, gen, path, stride, resident, spec, base](MTL::CommandBuffer*) { + // 阶段1（回调线程、持锁极短）：读惰性 id（此刻已算完）、预留侧区行。 + // 必须在回调里——id 只有 command buffer 完成后才有效。 + auto to_read = reserve(idp, n, layer, gen, resident, spec, base); + if (to_read.empty()) { + side_inflight_done(); // 无缺口可读：立即消账，避免 drain 空等 + return; + } + // 诊断门控 SIDEREGION_SYNC=1：回调内同步 pread+memcpy+publish（不派 bg）， + // 消除「异步 bg fill 与下一前向 gather 竞态」这一变量，用于 systematic-debugging 取证。 + static const bool kSync = []() { + const char* e = std::getenv("SIDEREGION_SYNC"); + return e && e[0] == '1'; + }(); + if (kSync) { + read_publish(ptrs, seg, to_read, path, stride, layer, gen); + side_inflight_done(); + return; + } + // 阶段2+3 派给自由后台线程：~40MB pread + memcpy + 发布 e2r 脱离 Metal 回调线程， + // 与主 stream 后续层计算、多层预取互相并发 → 真正恢复 I/O/计算重叠。 + // 闭包再次按值捕获 in：保活池/ids buffer 到后台读完（ptrs 指向其内存）。 + bg_submit_task([in, ptrs, seg, to_read, path, stride, layer, gen]() { + read_publish(ptrs, seg, to_read, path, stride, layer, gen); + side_inflight_done(); // 字节写完才消账 → drain 保证被消费行已就绪 + }); + }); + } + + private: + void run(const std::vector& in) { + std::vector ptrs; + ptrs.reserve(in.size() - 1); + for (size_t i = 1; i < in.size(); ++i) { + mx::array a = in[i]; + a.eval(); // 与 PrefetchStagingCachedPrimitive::eval_cpu 一致：先物化输入 + ptrs.push_back(a.data()); + } + mx::array ids = in[0]; + ids.eval(); + // CPU 路径同步执行（测试/无 GPU 时）：预留 + 读发布一气呵成。 + auto to_read = reserve(ids.data(), ids.size(), layer_, gen_, resident_, spec_, base_); + read_publish(ptrs, seg_, to_read, path_, stride_, layer_, gen_); + } + + // 阶段1：过滤常驻/去重 → 淘汰 ∉P 的旧行 → 为缺口预留物理行（出 free，暂不入 e2r， + // 避免消费者在字节写好前看到 e2r 命中而 gather 到脏字节）。返回 (expert, 预留行)。 + static std::vector> reserve( + const uint32_t* idp, size_t n, int layer, int gen, const std::vector& resident, + int spec, int base) { + const char* lfu_env = std::getenv("SIDEREGION_LFU"); // 每次读,便于测试切换 + // 默认开:持久 LFU 单缓冲=生产路径(cli 默认)。仅显式 SIDEREGION_LFU=0 回退旧 legacy 双缓冲。 + bool lfu = !lfu_env || lfu_env[0] != '0'; + std::unordered_set res(resident.begin(), resident.end()); + std::vector P; + std::unordered_set Pset, seen; + for (size_t i = 0; i < n; ++i) { + int e = static_cast(idp[i]); + if (res.count(e) || !seen.insert(e).second) continue; + P.push_back(e); + Pset.insert(e); + } + std::vector> to_read; + std::lock_guard lk(g_side_mutex); + SideLayer& c = g_side[{layer, gen}]; + if (!c.inited) { + for (int r = 0; r < spec; ++r) c.free_rows.push_back(base + r); + c.base = base; + c.spec = spec; + c.inited = true; + } + if (!lfu) { + // 旧行为:∉P 全弃(一次性预取批)。 + for (auto it = c.e2r.begin(); it != c.e2r.end();) { + if (!Pset.count(it->first)) { + c.free_rows.push_back(it->second); + it = c.e2r.erase(it); + } else { + ++it; + } + } + for (int e : P) { + if (c.e2r.count(e) || c.free_rows.empty()) continue; + to_read.emplace_back(e, c.free_rows.back()); + c.free_rows.pop_back(); + } + return to_read; + } + // LFU 持久:∉P 不清;再预测命中已驻专家 freq+1(越常预测越热)。 + for (int e : P) { + if (c.e2r.count(e)) c.freq[e] += 1; + } + for (int e : P) { + if (c.e2r.count(e)) continue; // 已驻,跳过(不重读) + int row; + if (!c.free_rows.empty()) { + row = c.free_rows.back(); + c.free_rows.pop_back(); + if (side_trace_hit(layer, row)) + fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d RESERVE_FROM_FREE row=%d expert=%d\n", + (unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, row, e); + } else { + // free 空:LFU 淘汰 e2r 中 freq 最小且 ∉P 者(tie-break:最小 expert id)。 + int victim = -1; + uint32_t best = 0; + for (auto& kv : c.e2r) { + if (Pset.count(kv.first)) continue; // 不淘本步要用的 + uint32_t f = c.freq.count(kv.first) ? c.freq[kv.first] : 0; + if (victim < 0 || f < best || (f == best && kv.first < victim)) { + victim = kv.first; + best = f; + } + } + if (victim < 0) continue; // 全是 P 热,无可淘 → 本步不读 + row = c.e2r[victim]; + c.e2r.erase(victim); + c.freq.erase(victim); + if (side_trace_hit(layer, row)) + fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d EVICT_REUSE row=%d victim=%d newExpert=%d\n", + (unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, row, victim, e); + } + to_read.emplace_back(e, row); + } + // ===== 临时诊断（SIDE_AUDIT=1）：reserve 结束时审计侧区行账本自洽性 ===== + if (std::getenv("SIDE_AUDIT")) side_audit(c, layer, gen, "reserve", to_read); + return to_read; + } + + // 审计：检测「一行被两专家占用」「free 与 e2r 行重叠」「to_read 行仍被 e2r 占用」。 + static void side_audit(SideLayer& c, int layer, int gen, const char* where, + const std::vector>& to_read) { + std::map row_owner; // row -> expert + for (auto& p : c.e2r) { + auto it = row_owner.find(p.second); + if (it != row_owner.end()) + fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d DOUBLE_OWN row=%d experts=%d,%d\n", + where, layer, gen, p.second, it->second, p.first); + row_owner[p.second] = p.first; + } + std::unordered_set free_set(c.free_rows.begin(), c.free_rows.end()); + for (auto& p : c.e2r) + if (free_set.count(p.second)) + fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d FREE_E2R_OVERLAP row=%d owned_by=%d\n", + where, layer, gen, p.second, p.first); + if (c.free_rows.size() != free_set.size()) + fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d FREE_DUP free_n=%zu uniq=%zu\n", + where, layer, gen, c.free_rows.size(), free_set.size()); + for (auto& pr : to_read) + if (row_owner.count(pr.second)) + fprintf(stderr, "[SIDE_AUDIT] %s L%d gen%d TOREAD_LIVE_ROW row=%d assigned_to=%d still_owned_by=%d\n", + where, layer, gen, pr.second, pr.first, row_owner[pr.second]); + } + + // 阶段2+3：pread blob 整行 + 各段 memcpy 直写进对应 per-key 池数组的物理侧区行（不持锁）， + // 写完后持锁发布 e2r。ptrs[i] 为第 i 个池 key 数组的 buffer 指针（C++ 拥有、地址恒定不迁移， + // 由 Route 3 底座保证——见 pool_owned_zeros），故后台异步直写安全、消费侧读同一 buffer。 + // 在后台线程跑：与计算并发，且不阻塞消费侧的 sideregion_kv。 + static void read_publish(const std::vector& ptrs, const std::vector& seg, + const std::vector>& to_read, + const std::string& path, size_t stride, int layer, int gen) { + // 段在 blob 记录内的偏移（段顺序 = seg 顺序 = ptrs/池 key 顺序）。 + std::vector seg_off(seg.size(), 0); + size_t acc = 0; + for (size_t i = 0; i < seg.size(); ++i) { seg_off[i] = acc; acc += static_cast(seg[i]); } + int fd = open_blob_nocache(path.c_str()); + if (fd < 0) { // 失败则把预留行还回 free + std::lock_guard lk(g_side_mutex); + SideLayer& c = g_side[{layer, gen}]; + for (auto& pr : to_read) c.free_rows.push_back(pr.second); + return; + } + std::vector rec(stride); // 整条 blob 记录临时缓冲 + std::vector> done; + for (auto& pr : to_read) { + int e = pr.first, row = pr.second; + if (::pread(fd, rec.data(), stride, static_cast(static_cast(e) * stride)) != + static_cast(stride)) { + std::lock_guard lk(g_side_mutex); // 读失败：行还回 free + g_side[{layer, gen}].free_rows.push_back(row); + continue; + } + // 各段 memcpy 进对应 per-key 池数组的物理行 row（(n_slots,*shape) 布局 → 行偏移 = row*seg[i]）。 + for (size_t i = 0; i < seg.size(); ++i) + std::memcpy(ptrs[i] + static_cast(row) * static_cast(seg[i]), + rec.data() + seg_off[i], static_cast(seg[i])); + if (side_trace_hit(layer, row)) + fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d WRITEPOOL row=%d expert=%d\n", + (unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, row, e); + done.emplace_back(e, row); + } + ::close(fd); + { + std::lock_guard lk(g_side_mutex); + SideLayer& c = g_side[{layer, gen}]; + for (auto& pr : done) { // 字节就绪后才发布 e2r + c.e2r[pr.first] = pr.second; + if (!c.freq.count(pr.first)) c.freq[pr.first] = 1; // 新专家初始 freq + if (side_trace_hit(layer, pr.second)) + fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d PUBLISH row=%d expert=%d\n", + (unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, pr.second, + pr.first); + } + if (std::getenv("SIDE_AUDIT")) side_audit(c, layer, gen, "publish", {}); + } + } + std::vector seg_; + int layer_; + int gen_; + std::string path_; + size_t stride_; + std::vector resident_; + int spec_; + int base_; +}; + +mx::array prefetch_pool_sideregion( + const std::vector& pool_list, const std::vector& seg_nbytes, + const mx::array& expert_ids, int layer, const std::string& path, int stride, + const std::vector& resident, int spec_slots, int base_row, int gen, + mx::StreamOrDevice s = {}) { + // 防御：段数必须与池数组个数一一对应（每段映射唯一一个 per-key 数组）。 + if (seg_nbytes.size() != pool_list.size()) { + throw std::invalid_argument( + "prefetch_pool_sideregion: seg_nbytes.size()=" + std::to_string(seg_nbytes.size()) + + " != pool_list.size()=" + std::to_string(pool_list.size())); + } + // 防御：各段字节之和必须等于 stride（blob 记录恰好是各段的拼接）。 + size_t seg_sum = 0; + for (int b : seg_nbytes) seg_sum += static_cast(b); + if (seg_sum != static_cast(stride)) { + throw std::invalid_argument( + "prefetch_pool_sideregion: sum(seg_nbytes)=" + std::to_string(seg_sum) + + " != stride=" + std::to_string(stride)); + } + std::vector inputs; + inputs.push_back(expert_ids); + for (auto& a : pool_list) inputs.push_back(a); + return mx::array( + mx::Shape{1}, mx::uint8, + std::make_shared( + mx::to_stream(s), seg_nbytes, layer, gen, path, static_cast(stride), resident, + spec_slots, base_row), + inputs); +} + +std::vector sideregion_contents(int layer, int gen) { + std::lock_guard lk(g_side_mutex); + std::vector out; + auto it = g_side.find({layer, gen}); + if (it != g_side.end()) + for (auto& p : it->second.e2r) { out.push_back(p.first); out.push_back(p.second); } + return out; +} + +// 侧区 e2r → 两个 device mx.array (keys uint32, vals int32),直接在 C++ 从 map 建连续 buffer。 +// 消掉 Python 侧 dict 构建 + list(...)→mx.array 的每层 host 胶水。 +std::pair sideregion_kv(int layer, int gen) { + std::vector keys; + std::vector vals; + { + std::lock_guard lk(g_side_mutex); + auto it = g_side.find({layer, gen}); + if (it != g_side.end()) { + keys.reserve(it->second.e2r.size()); + vals.reserve(it->second.e2r.size()); + for (auto& p : it->second.e2r) { + keys.push_back(static_cast(p.first)); + vals.push_back(static_cast(p.second)); + if (side_trace_hit(layer, p.second)) + fprintf(stderr, "[SIDE_TRACE ev=%llu tid=%u] L%d gen%d CONSUME_KV row=%d expert=%d\n", + (unsigned long long)g_side_ev.fetch_add(1), side_tid(), layer, gen, p.second, p.first); + } + } + } + int n = static_cast(keys.size()); + return {mx::array(keys.data(), mx::Shape{n}, mx::uint32), + mx::array(vals.data(), mx::Shape{n}, mx::int32)}; +} + +void sideregion_reset() { + std::lock_guard lk(g_side_mutex); + g_side.clear(); +} + +std::unordered_map sideregion_snapshot(int layer, int gen) { + std::unordered_map side; + std::lock_guard lk(g_side_mutex); + auto it = g_side.find({layer, gen}); + if (it != g_side.end()) for (auto& p : it->second.e2r) side[p.first] = p.second; + return side; +} diff --git a/native/ext/pool/side_region.h b/native/ext/pool/side_region.h new file mode 100644 index 0000000..209935b --- /dev/null +++ b/native/ext/pool/side_region.h @@ -0,0 +1,21 @@ +// 段散写持久侧区缓存（zero-copy dual-source 默认路径）：keep/evict/read-delta， +// 命中缺口时 pread blob 行并把各段 memcpy 进对应 per-key 池数组的 (base_row+cache_row) 行。 +#pragma once +#include "../common.h" +#include + +mx::array prefetch_pool_sideregion( + const std::vector& pool_list, const std::vector& seg_nbytes, + const mx::array& expert_ids, int layer, const std::string& path, int stride, + const std::vector& resident, int spec_slots, int base_row, int gen, + mx::StreamOrDevice s); +std::vector sideregion_contents(int layer, int gen); // [expert, phys_row, ...] +std::pair sideregion_kv(int layer, int gen); // (keys uint32, vals int32) device 数组 +void sideregion_reset(); +// 排空侧区在途 fill：阻塞直到所有已提交的侧区预取字节写完。消费方在前向开头调用， +// 保证被消费的侧区行字节已完全写好（闭合异步写-读竞态）。 +void sideregion_drain(); + +// 内部接口（不绑 Python）：取某层某代侧区 e2r 快照，供 demand_dual 在自己的核心状态机里 +// 叠加侧区命中。等价于在 g_side 锁下拷贝 e2r（与原 demand_dual 内联循环一字不差）。 +std::unordered_map sideregion_snapshot(int layer, int gen); diff --git a/native/ext/prefetch/prefetch.cpp b/native/ext/prefetch/prefetch.cpp new file mode 100644 index 0000000..883addf --- /dev/null +++ b/native/ext/prefetch/prefetch.cpp @@ -0,0 +1,217 @@ +// [2] 轻量预取(无 staging，仅预热 page cache) + [3] staging 版 miss→hit + STAGING_HPROF 探针。 +#include "prefetch.h" +#include "../io/bg_reader.h" + +#include +#include +#include +#include +#include +#include +#include + +// ====== [2] 轻量预取(无 staging 模式)：GPU 完成回调里读 inds、pread 预热 page cache ====== +// 命门验证:回调在 command buffer 完成后触发,此时 inds 已算完,读到的是正确值。 +class PrefetchOnCompletePrimitive : public mx::Primitive { + public: + PrefetchOnCompletePrimitive(mx::Stream s, std::string path, size_t stride, bool do_read) + : Primitive(s), path_(std::move(path)), stride_(stride), do_read_(do_read) {} + const char* name() const override { return "PrefetchOnCompletePrimitive"; } + + void eval_cpu(const std::vector& inputs, std::vector& outputs) override { + outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes())); + mx::array ids = inputs[0]; + ids.eval(); + record_and_read(ids.data(), ids.size(), path_, stride_, do_read_); + } + + void eval_gpu(const std::vector& inputs, std::vector& outputs) override { + outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes())); + mx::array ids = inputs[0]; // by-value：保活到回调结束 + const uint32_t* ptr = ids.data(); // 指针此刻有效（buffer 已分配），值待 GPU 算完 + size_t n = ids.size(); + std::string path = path_; + size_t stride = stride_; + bool do_read = do_read_; + auto& enc = mx::metal::get_command_encoder(stream()); + MTL::CommandBuffer* cb = enc.get_command_buffer(); + cb->addCompletedHandler([ids, ptr, n, path, stride, do_read](MTL::CommandBuffer*) { + // 此时 buffer 已完成 → ptr 指向已算好的 inds 值。ids 按值捕获保证 buffer 不被释放。 + record_and_read(ptr, n, path, stride, do_read); + }); + } + + private: + static void record_and_read(const uint32_t* p, size_t n, const std::string& path, + size_t stride, bool do_read) { + int fd = (do_read && !path.empty()) ? ::open(path.c_str(), O_RDONLY) : -1; + static thread_local std::vector buf; + if (do_read && buf.size() < stride) buf.resize(stride); + for (size_t i = 0; i < n; ++i) { + int e = static_cast(p[i]); + if (fd >= 0) ::pread(fd, buf.data(), stride, static_cast(static_cast(e) * stride)); + } + if (fd >= 0) ::close(fd); + } + std::string path_; + size_t stride_; + bool do_read_; +}; + +// 返回一个 dummy 输出（挂进图里，eval 时触发上面的回调）。 +mx::array prefetch_on_complete( + const mx::array& expert_ids, + const std::string& path, + int stride, + bool do_read = true, + mx::StreamOrDevice s = {}) { + return mx::array( + mx::Shape{1}, + mx::uint8, + std::make_shared( + mx::to_stream(s), path, static_cast(stride), do_read), + std::vector{expert_ids}); +} + +// ---- handler 触发时刻探针(STAGING_HPROF)：记录每次 staging 完成回调被 Metal 触发的时刻 ---- +// 用于实测"完成回调到底扎堆在 eval 尾、还是逐层铺开"。仅诊断用，默认关。 +static std::mutex g_hprof_mutex; +static std::vector> g_hprof_log; // (gen, layer, t_fire_seconds) +static bool g_hprof_on = false; + +static inline double hprof_steady_now() { + return std::chrono::duration( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); +} + +// ====== [3] staging 版 miss→hit：handler pread 进 per-layer staging + 记录 (expert→row) ====== +static std::mutex g_stg_mutex; +// layer -> (gen, [(expert, row)])；handler 原子写 gen+映射，主线程按 gen 匹配 buffer 后 take。 +static std::map>>> g_stg_ready; + +class PrefetchStagingPrimitive : public mx::Primitive { + public: + // 方案B：expert_ids 是"预测宽集合"(top-N，按门控分降序)；resident 是目标层当前常驻专家。 + // handler 在回调里过滤掉常驻、按降序取前 cap 个缺口 pread 进 staging（cap=buffer 行数）。 + // 这样预测可以很宽(高 recall)，而 staging 内存只按 cap 预留(小，覆盖缺口分布即可)。 + PrefetchStagingPrimitive(mx::Stream s, int layer, long gen, std::string path, size_t stride, + std::vector resident, int cap, bool parallel) + : Primitive(s), layer_(layer), gen_(gen), path_(std::move(path)), stride_(stride), + resident_(std::move(resident)), cap_(cap), parallel_(parallel) {} + const char* name() const override { return "PrefetchStagingPrimitive"; } + + void eval_cpu(const std::vector& inputs, std::vector& outputs) override { + outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes())); + mx::array ids = inputs[0]; ids.eval(); + mx::array stg = inputs[1]; stg.eval(); + fill(ids.data(), ids.size(), stg.data(), layer_, gen_, path_, stride_, + resident_, cap_); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) override { + outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes())); + mx::array ids = inputs[0]; + mx::array stg = inputs[1]; + const uint32_t* idp = ids.data(); + uint8_t* sp = stg.data(); + size_t n = ids.size(); + int layer = layer_; long gen = gen_; std::string path = path_; size_t stride = stride_; + std::vector resident = resident_; int cap = cap_; bool parallel = parallel_; + auto& enc = mx::metal::get_command_encoder(stream()); + MTL::CommandBuffer* cb = enc.get_command_buffer(); + cb->addCompletedHandler( + [ids, stg, idp, sp, n, layer, gen, path, stride, resident, cap, parallel](MTL::CommandBuffer*) { + // 探针：记录本回调被 Metal 触发的时刻(= 该层 pread 能开始跑的时刻)。 + if (g_hprof_on) { + double t = hprof_steady_now(); + std::lock_guard lk(g_hprof_mutex); + g_hprof_log.emplace_back(gen, layer, t); + } + // ids/stg 按值捕获保活 buffer；派后台线程时再拷一份保活到 fill 跑完。 + if (parallel) { + bg_submit_task([ids, stg, idp, sp, n, layer, gen, path, stride, resident, cap]() { + fill(idp, n, sp, layer, gen, path, stride, resident, cap); + }); + } else { + fill(idp, n, sp, layer, gen, path, stride, resident, cap); + } + }); + } + + private: + static void fill(const uint32_t* idp, size_t n, uint8_t* stg, + int layer, long gen, const std::string& path, size_t stride, + const std::vector& resident, int cap) { + int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) return; + std::unordered_set res(resident.begin(), resident.end()); + std::unordered_set done; // 去重：同一缺口只 pread 一次 + std::vector> ready; + ready.reserve(static_cast(cap)); + int row = 0; // 写入的 staging 行（≤ cap） + for (size_t i = 0; i < n && row < cap; ++i) { + int e = static_cast(idp[i]); + if (res.count(e) || !done.insert(e).second) continue; // 已常驻/已取过 → 跳过 + if (::pread(fd, stg + static_cast(row) * stride, stride, + static_cast(static_cast(e) * stride)) + == static_cast(stride)) { + ready.emplace_back(e, row); + ++row; + } + } + ::close(fd); + std::lock_guard lk(g_stg_mutex); + g_stg_ready[layer] = {gen, std::move(ready)}; // 原子：gen 与映射一起写 + } + int layer_; + long gen_; + std::string path_; + size_t stride_; + std::vector resident_; // 目标层提交时刻的常驻专家快照（过滤用） + int cap_; // staging buffer 行数上限 + bool parallel_; // true: fill 派后台线程池并行；false: 回调线程同步(旧行为) +}; + +mx::array prefetch_into_staging( + const mx::array& staging, const mx::array& expert_ids, int layer, long gen, + const std::string& path, int stride, const std::vector& resident, int cap, + bool parallel, mx::StreamOrDevice s = {}) { + return mx::array( + mx::Shape{1}, mx::uint8, + std::make_shared( + mx::to_stream(s), layer, gen, path, static_cast(stride), resident, cap, parallel), + std::vector{expert_ids, staging}); +} + +// 取走某层就绪记录：[gen, e0,r0,e1,r1,...]（首元素是 generation）；空表示无就绪。 +std::vector prefetch_staging_take(int layer) { + std::lock_guard lk(g_stg_mutex); + std::vector out; + auto it = g_stg_ready.find(layer); + if (it != g_stg_ready.end()) { + out.push_back(it->second.first); // gen + for (auto& p : it->second.second) { out.push_back(p.first); out.push_back(p.second); } + g_stg_ready.erase(it); + } + return out; +} + +// ---- handler 触发时刻探针接口 ---- +void staging_hprof_enable(bool on) { + std::lock_guard lk(g_hprof_mutex); + g_hprof_on = on; + g_hprof_log.clear(); // 开启即清零，便于每次采集干净 +} +double staging_hprof_now() { return hprof_steady_now(); } // 与日志同一时钟，供 Python 标 eval 边界 +// 扁平返回 [gen0,layer0,t0, gen1,layer1,t1, ...]（避开 nanobind tuple caster）。 +std::vector staging_hprof_get() { + std::lock_guard lk(g_hprof_mutex); + std::vector out; + out.reserve(g_hprof_log.size() * 3); + for (auto& r : g_hprof_log) { + out.push_back(static_cast(std::get<0>(r))); + out.push_back(static_cast(std::get<1>(r))); + out.push_back(std::get<2>(r)); + } + return out; +} diff --git a/native/ext/prefetch/prefetch.h b/native/ext/prefetch/prefetch.h new file mode 100644 index 0000000..2d9cc2a --- /dev/null +++ b/native/ext/prefetch/prefetch.h @@ -0,0 +1,23 @@ +// GPU 完成回调预取：算完 inds 后在回调里读 id、pread 预热 page cache 或写 per-layer staging。 +#pragma once +#include "../common.h" + +// 挂 GPU 完成回调：算完 inds 后在回调里读 id 预热 page cache，返回 dummy。 +// （无 staging 的轻量预取模式：仅把预测专家字节读进 page cache，不落池。） +mx::array prefetch_on_complete( + const mx::array& expert_ids, const std::string& path, int stride, + bool do_read, mx::StreamOrDevice s); + +// miss→hit（方案B）：expert_ids 为预测宽集合(降序)，回调按 resident 过滤、取前 cap 个缺口 +// pread 进 per-layer staging buffer(cap=buffer 行数)，并原子记录 (gen,[expert,row])。 +mx::array prefetch_into_staging( + const mx::array& staging, const mx::array& expert_ids, int layer, long gen, + const std::string& path, int stride, const std::vector& resident, int cap, + bool parallel, mx::StreamOrDevice s); +// 取走某层就绪记录：[gen, e0,r0,e1,r1,...]；空表示无就绪。 +std::vector prefetch_staging_take(int layer); + +// handler 触发时刻探针(诊断用)：enable 开关并清零；now 取同时钟当前秒；get 取 (gen,layer,t) 日志。 +void staging_hprof_enable(bool on); +double staging_hprof_now(); +std::vector staging_hprof_get(); // 扁平 [gen,layer,t, ...] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5a7aa08 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,44 @@ +[project] +name = "sparkle-engine" +version = "0.1.0" +description = "Sparkle：Apple Silicon 上流式 MoE 推理（Qwen3-Next-80B 8bit）" +readme = "README.md" +# 仓库自带 native_moe_ext*.cpython-314*.so，须与扩展 ABI 一致 +requires-python = "==3.14.*" +dependencies = [ + "fastapi>=0.115", + "mlx>=0.31", + "mlx-lm>=0.31", + "numpy>=2.0", + "textual>=0.80", + "uvicorn>=0.30", +] + +[project.scripts] +sparkle = "mlx_streaming.cli:main" + +# 仅重编 native / 跑单测时需要；日常推理用已有 .venv + .so 即可 +[dependency-groups] +dev = [ + "httpx>=0.27", + "nanobind>=2.12.0", + "pytest>=8", + "pytest-asyncio>=0.23", +] + +[[tool.uv.index]] +# 清华 PyPI 镜像 +url = "https://pypi.tuna.tsinghua.edu.cn/simple" +default = true + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["mlx_streaming"] + +[tool.pytest.ini_options] +addopts = "-q" +testpaths = ["mlx_streaming/tests"] +asyncio_mode = "auto" diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..5ba4be0 --- /dev/null +++ b/uv.lock @@ -0,0 +1,1129 @@ +version = 1 +revision = 3 +requires-python = ">=3.11" + +[[package]] +name = "annotated-doc" +version = "0.0.4" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, +] + +[[package]] +name = "annotated-types" +version = "0.8.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, +] + +[[package]] +name = "anyio" +version = "4.13.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "idna" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, +] + +[[package]] +name = "certifi" +version = "2026.5.20" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f3/ce/ee2ecad540810a79593028e88299baeae54d346cc7a0d94b6199988b89b1/certifi-2026.5.20.tar.gz", hash = "sha256:69dea482ab64caa7b9f6aba1c6bf48bb6a5448d1c0f1b17ab42ad8c763a5344d", size = 135422, upload-time = "2026-05-20T11:46:50.073Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/59/8c/57e832b7af6d7c5abe66eb3fbe3a3a32f4d11ea23a1aa7131371035be991/certifi-2026.5.20-py3-none-any.whl", hash = "sha256:3c52e209ba0a4ad7aebe60436a4ab349c39e1e602e8c134221e546902ad25897", size = 134134, upload-time = "2026-05-20T11:46:48.578Z" }, +] + +[[package]] +name = "click" +version = "8.4.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9b/98/518d8e5081007684232226f475082b30087d0f585e8457db087298259f49/click-8.4.1.tar.gz", hash = "sha256:918b5633eddf6b41c32d4f454bf0de810065c74e3f7dbf8ee5452f8be88d3e96", size = 353007, upload-time = "2026-05-22T04:08:37.769Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "fastapi" +version = "0.140.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "pydantic" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0d/fb/fd7671137d9fa3df1d93a2f5111eb982709201724b29f211e4beb2d58688/fastapi-0.140.0.tar.gz", hash = "sha256:f338951b82fd74ca8f843163aec43ea1a1ce84d515415a50fa98fa25572a5544", size = 420968, upload-time = "2026-07-24T21:16:41.187Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/eb/76/6d9e25ad88da9d3ff744bcdbec4736e38c2288611d43f673a5d9bfa27c07/fastapi-0.140.0-py3-none-any.whl", hash = "sha256:e951c0a0d9540bf5d9a2a9e078fd415da2ab7e312d435139e7d9e2e7fe9f0b23", size = 130863, upload-time = "2026-07-24T21:16:42.89Z" }, +] + +[[package]] +name = "filelock" +version = "3.29.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1f/f9/f38573ed5844586db374d085911740a501ccfa373b455fc9413f09f85237/filelock-3.29.1.tar.gz", hash = "sha256:d97e6b1b9757569626c58caa07dc4beb1613f4a2938b1e8cc81afca398906c9e", size = 59335, upload-time = "2026-06-03T15:19:04.053Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4c/a0/614c5fe402fd88951df45f4dda2fa3b4e17a99ecd92340771929169b3b95/filelock-3.29.1-py3-none-any.whl", hash = "sha256:85199dfd706869641b72b2e8955d5416a4b2b7dc4b0e8e6d97b4cc1299a6983b", size = 40750, upload-time = "2026-06-03T15:19:02.959Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.4.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d5/8d/1c51c094345df128ca4a990d633fe1a0ff28726c9e6b3c41ba65087bba1d/fsspec-2026.4.0.tar.gz", hash = "sha256:301d8ac70ae90ef3ad05dcf94d6c3754a097f9b5fe4667d2787aa359ec7df7e4", size = 312760, upload-time = "2026-04-29T20:42:38.635Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d5/0c/043d5e551459da400957a1395e0febbf771446ff34291afcbe3d8be2a279/fsspec-2026.4.0-py3-none-any.whl", hash = "sha256:11ef7bb35dab8a394fde6e608221d5cf3e8499401c249bebaeaad760a1a8dec2", size = 203402, upload-time = "2026-04-29T20:42:36.842Z" }, +] + +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.5.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/74/d8/5c06fc76461418326a7decf8367480c35be11a41fd938633929c60a9ec6b/hf_xet-1.5.0.tar.gz", hash = "sha256:e0fb0a34d9f406eed88233e829a67ec016bec5af19e480eac65a233ea289a948", size = 837196, upload-time = "2026-05-06T06:18:15.583Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/68/9b/6912c99070915a4f28119e3c5b52a9abd1eec0ad5cb293b8c967a0c6f5a2/hf_xet-1.5.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:7d70fe2ce97b9db73b9c9b9c81fe3693640aec83416a966c446afea54acfae3c", size = 4023383, upload-time = "2026-05-06T06:17:53.947Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0f/6d/9563cfde59b5d8128a9c7ec972a087f4c782e4f7bac5a85234edfd5d5e49/hf_xet-1.5.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:73a0dae8c71de3b0633a45c73f4a4a5ed09e94b43441d82981a781d4f12baa42", size = 3792751, upload-time = "2026-05-06T06:17:51.791Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/07/a5/ed5a0cf35b49a0571af5a8f53416dad1877a718c021c9937c3a53cb45781/hf_xet-1.5.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a60290ec57e9b71767fba7c3645ddafdd0759974b540441510c629c6db6db24a", size = 4456058, upload-time = "2026-05-06T06:17:40.735Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/fb/3ae8bf2a7a37a4197d0195d7247fd25b3952e15cb8a599e285dfaa6f52b3/hf_xet-1.5.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e5de0f6deada0dada870bb376a11bcd1f08abf3a968a6d118f33e72d1b1eb480", size = 4250783, upload-time = "2026-05-06T06:17:38.412Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a2/9b/8bae40d4d91525085137196e84eb0ed49cf65b5e96e5c3ecdadd8bd0fac2/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c799d49f1a5544a0ef7591c0ee75e0d6b93d6f56dc7a4979f59f7518d2872216", size = 4445594, upload-time = "2026-05-06T06:18:04.219Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/13/59/c74efbbd4e8728172b2cc72a2bc014d2947a4b7bdced932fbd3f5da1a4e5/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2baea1b0b989e5c152fe81425f7745ddc8901280ba3d97c98d8cdece7b706c60", size = 4663995, upload-time = "2026-05-06T06:18:06.1Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/32/8e1e0410af64cda9b139d1dcebdc993a8ff9c8c7c0e2696ae356d75ccc0d/hf_xet-1.5.0-cp313-cp313t-win_amd64.whl", hash = "sha256:526345b3ed45f374f6317349df489167606736c876241ba984105afe7fd4839d", size = 3966608, upload-time = "2026-05-06T06:18:19.74Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fc/34/a8febc8f4edbea8b3e21b02ebc8b628679b84ba7e45cde624a7736b51500/hf_xet-1.5.0-cp313-cp313t-win_arm64.whl", hash = "sha256:786d28e2eb8315d5035544b9d137b4a842d600c434bb91bf7d0d953cce906ad4", size = 3796946, upload-time = "2026-05-06T06:18:17.568Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2a/20/8fc8996afe5815fa1a6be8e9e5c02f24500f409d599e905800d498a4e14d/hf_xet-1.5.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:872d5601e6deea30d15865ede55d29eac6daf5a534ab417b99b6ef6b076dd96c", size = 4023495, upload-time = "2026-05-06T06:18:01.94Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/32/6a/93d84463c00cecb561a7508aa6303e35ee2894294eac14245526924415fe/hf_xet-1.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:9929561f5abf4581c8ea79587881dfef6b8abb2a0d8a51915936fc2a614f4e73", size = 3792731, upload-time = "2026-05-06T06:18:00.021Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9d/5a/8ec8e0c863b382d00b3c2e2af6ded6b06371be617144a625903a6d562f4b/hf_xet-1.5.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f7b7bbae318e583a86fb21e5a4a175d6721d628a2874f4bd022d0e660c32a682", size = 4456738, upload-time = "2026-05-06T06:17:49.574Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c5/ca/f7effa1a67717da2bcc6b6c28f71c6ca648c77acaec4e2c32f40cbe16d85/hf_xet-1.5.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:cf7b2dc6f31a4ea754bb50f74cde482dcf5d366d184076d8530b9872787f3761", size = 4251622, upload-time = "2026-05-06T06:17:47.096Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/65/f2/19247dba3e231cf77dec59ddfb878f00057635ff773d099c9b59d37812c3/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8dbcbab554c9ef158ef2c991545c3e970ddd8cc7acdcd0a78c5a41095dab4ded", size = 4445667, upload-time = "2026-05-06T06:18:11.983Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7f/64/6f116801a3bcfb6f59f5c251f48cadc47ea54026441c4a385079286a94fa/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5906bf7718d3636dc13402914736abe723492cb730f744834f5f5b67d3a12702", size = 4664619, upload-time = "2026-05-06T06:18:13.771Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5c/e8/069542d37946ed08669b127e1496fa99e78196d71de8d41eda5e9f1b7a58/hf_xet-1.5.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f3dc2248fc01cc0a00cd392ab497f1ca373fcbc7e3f2da1f452480b384e839e", size = 3966802, upload-time = "2026-05-06T06:18:28.162Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f9/91/fc6fdec27b14d04e88c386ac0a0129732b53fa23f7c4a78f4b83a039c567/hf_xet-1.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:b285cea1b5bab46b758772716ba8d6854a1a0310fed1c249d678a8b38601e5a0", size = 3797168, upload-time = "2026-05-06T06:18:26.287Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3d/fb/69ff198a82cae7eb1a69fb84d93b3a3e4816564d76817fe541ddc96874eb/hf_xet-1.5.0-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:dad0dc84e941b8ba3c860659fe1fdc35c049d47cce293f003287757e971a8f56", size = 4030814, upload-time = "2026-05-06T06:17:57.933Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9b/ff/edcc2b40162bef3ff78e14ab637e5f3b89243d6aee72f5949d3bb6a5af83/hf_xet-1.5.0-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:fd6e5a9b0fdac4ed03ed45ef79254a655b1aaab514a02202617fbf643f5fdf7a", size = 3798444, upload-time = "2026-05-06T06:17:55.79Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/49/4d/103f76b04310e5e57656696cc184690d20c466af0bca3ca88f8c8ea5d4f3/hf_xet-1.5.0-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3531b1823a0e6d77d80f9ed15ca0e00f0d115094f8ac033d5cae88f4564cc949", size = 4465986, upload-time = "2026-05-06T06:17:44.886Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c4/a2/546f47f464737b3edbab6f8ddb57f2599b93d2cbb66f06abb475ccb48651/hf_xet-1.5.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9a0ee58cd18d5ea799f7ed11290bbccbe56bdd8b1d97ca74b9cc49a3945d7a3b", size = 4259865, upload-time = "2026-05-06T06:17:42.639Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/95/7f/1be593c1f28613be2e196473481cd81bfc5910795e30a34e8f744f6cac4f/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1e60df5a42e9bed8628b6416af2cba4cba57ae9f02de226a06b020d98e1aab18", size = 4459835, upload-time = "2026-05-06T06:18:08.026Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/aa/b2/703569fc881f3284487e68cda7b42179978480da3c438042a6bbbb4a671c/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4b35549ce62601b84da4ff9b24d970032ace3d4430f52d91bcbb26c901d6c690", size = 4672414, upload-time = "2026-05-06T06:18:09.864Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/af/37/1b6def445c567286b50aa3b33828158e135b1be44938dde59f11382a500c/hf_xet-1.5.0-cp37-abi3-win_amd64.whl", hash = "sha256:2806c7c17b4d23f8d88f7c4814f838c3b6150773fe339c20af23e1cfaf2797e4", size = 3977238, upload-time = "2026-05-06T06:18:23.621Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/62/94/3b66b148778ee100dcfd69c2ca22b57b41b44d3063ceec934f209e9184ce/hf_xet-1.5.0-cp37-abi3-win_arm64.whl", hash = "sha256:b6c9df403040248c76d808d3e047d64db2d923bae593eb244c41e425cf6cd7be", size = 3806916, upload-time = "2026-05-06T06:18:21.7Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.18.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/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'" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "tqdm" }, + { name = "typer" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/d8/748ea0a47f0fa15227fe682f7a80826b4b7c096e4818044b8f56d6cb66d6/huggingface_hub-1.18.0.tar.gz", hash = "sha256:f0c5ecd1ef8c6a60f86f61ee278f2c1570ba9e279c9f54de9094210723b3613b", size = 812699, upload-time = "2026-06-05T09:26:33.401Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0b/03/40a05316cb6616e5b7efd7773656441ab04b4b022c2199e79bb4622a92a3/huggingface_hub-1.18.0-py3-none-any.whl", hash = "sha256:729be4a976fb706dcc02d176bcda8a3f32bdf21a294e8f4b3dda6fbcbc9c1ab1", size = 684411, upload-time = "2026-06-05T09:26:31.48Z" }, +] + +[[package]] +name = "idna" +version = "3.18" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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 = "linkify-it-py" +version = "2.1.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "uc-micro-py" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2e/c9/06ea13676ef354f0af6169587ae292d3e2406e212876a413bf9eece4eb23/linkify_it_py-2.1.0.tar.gz", hash = "sha256:43360231720999c10e9328dc3691160e27a718e280673d444c38d7d3aaa3b98b", size = 29158, upload-time = "2026-03-01T07:48:47.683Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b4/de/88b3be5c31b22333b3ca2f6ff1de4e863d8fe45aaea7485f591970ec1d3e/linkify_it_py-2.1.0-py3-none-any.whl", hash = "sha256:0d252c1594ecba2ecedc444053db5d3a9b7ec1b0dd929c8f1d74dce89f86c05e", size = 19878, upload-time = "2026-03-01T07:48:46.098Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.optional-dependencies] +linkify = [ + { name = "linkify-it-py" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/08/db/fefacb2136439fc8dd20e797950e749aa1f4997ed584c62cfb8ef7c2be0e/markupsafe-3.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad", size = 11631, upload-time = "2025-09-27T18:36:18.185Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/0f/62/d9c46a7f5c9adbeeeda52f5b8d802e1094e9717705a645efc71b0913a0a8/markupsafe-3.0.3-cp311-cp311-win32.whl", hash = "sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19", size = 14572, upload-time = "2025-09-27T18:36:28.045Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/83/8a/4414c03d3f891739326e1783338e48fb49781cc915b2e0ee052aa490d586/markupsafe-3.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01", size = 15077, upload-time = "2025-09-27T18:36:29.025Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/35/73/893072b42e6862f319b5207adc9ae06070f095b358655f077f69a35601f0/markupsafe-3.0.3-cp311-cp311-win_arm64.whl", hash = "sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c", size = 13876, upload-time = "2025-09-27T18:36:29.954Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + +[[package]] +name = "mdit-py-plugins" +version = "0.6.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "markdown-it-py" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/59/fc/f8d0863f8862f25602c0404d75568e89fb6b4109804645e5cdfb1be5cf56/mdit_py_plugins-0.6.1.tar.gz", hash = "sha256:a2bca0f039f39dbd35fb74ae1b5f998608c437463371f0ff7f49a19a17a114d0", size = 56114, upload-time = "2026-05-13T09:03:38.91Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a5/69/6da5581c6a7fede7dc261bf4e67d6adca4196f176b43288b55b3db395b6e/mdit_py_plugins-0.6.1-py3-none-any.whl", hash = "sha256:214c82fb2ac524472ab6a5bcab1de80f73b50443e187f401bfd77efbc7c6481d", size = 66663, upload-time = "2026-05-13T09:03:37.76Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.31.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, +] +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/94/89/1e77ec3ff380e8fb9e7258047374d31452a0f9828a0e370f127b07dd8288/mlx-0.31.2-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:4a3f181b367d404e44a6bd68ef5eb573930809ac60cacd51d0c851c629b1b651", size = 586911, upload-time = "2026-04-22T03:14:29.675Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6a/41/c1907f05f8a3fc54025fb78ad68d3c4a4b931664d03c0a24f7f431cc4087/mlx-0.31.2-cp311-cp311-macosx_15_0_arm64.whl", hash = "sha256:70297cbef7479429f69c966bfed10da20a6f0c2aa997eec2b4f6ba1a07caf2ef", size = 586915, upload-time = "2026-04-22T03:14:31.403Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/97/b0/61ac2c14773c786fecbda28067b0207a0c654cb4d10c548808c51284d700/mlx-0.31.2-cp311-cp311-macosx_26_0_arm64.whl", hash = "sha256:c0ff158b7ac93a4b5659adbc70053498b30a5964fc45f78596398e056a96c36a", size = 587030, upload-time = "2026-04-22T03:14:32.961Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/de/53/e12feb7078ee472983555fcb1da4749a2bbbc8fc5b29b78c205b96d37d1e/mlx-0.31.2-cp311-cp311-manylinux_2_35_aarch64.whl", hash = "sha256:cd5d42b0b2bee7efe1b0680a7e302943dd33b92c879cffa0358ffdb5a4a8d27b", size = 652994, upload-time = "2026-04-22T03:14:34.691Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c5/40/f92c8cdc9595bf24c7e483a3156bfe0cc99a5cf5545d8dba8e7fe000c10b/mlx-0.31.2-cp311-cp311-manylinux_2_35_x86_64.whl", hash = "sha256:b368f7ede4238cc44076e4843820338c453c21ee50bd3ee26d4b182c179fd8e1", size = 692086, upload-time = "2026-04-22T03:14:36.45Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c3/47/5f33906cb03d6a378a697cd2d2641a26b37dea17ee3d9124d7e39e8eca01/mlx-0.31.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:e5067aaf2be1f3d7bba5be52348775804f111173c1ed04639618fd713b1a530f", size = 584863, upload-time = "2026-04-22T03:14:38.211Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/08/e7/a851a451b1327af9fb4df3991b9ae87d066b6f6630e854af55c288b0995a/mlx-0.31.2-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:edb9797db7d852477ca1c99708058654ee860d4148fe5765f0d55528e2b1aa22", size = 584860, upload-time = "2026-04-22T03:14:39.746Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3b/15/0d1dc0597644e5e7b011ca954ba0c47e13cd880a3b909b0c3f1b4d8bf8f1/mlx-0.31.2-cp312-cp312-macosx_26_0_arm64.whl", hash = "sha256:51ca102db641b01e7cb083ce8ecb580e281530a141a7ca12544bb370641630ae", size = 584887, upload-time = "2026-04-22T03:14:41.585Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5d/c3/00664239a98e8bd614733c4182cd402d2bacad2d7f79eca66562ac406870/mlx-0.31.2-cp312-cp312-manylinux_2_35_aarch64.whl", hash = "sha256:117c7583cae0ca107cd53c591cc34f8e75f97a505aa47088844b7dc0fc69dc67", size = 627863, upload-time = "2026-04-22T03:14:43.326Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/53/7b/af6cd73a79772af6f19eab2cb4c48eda23a9294d1650a4c1269a9996e532/mlx-0.31.2-cp312-cp312-manylinux_2_35_x86_64.whl", hash = "sha256:99572133181481640a8bf8d449daf083816d0af3ee050c8adfc5bf45ceca91c6", size = 685090, upload-time = "2026-04-22T03:14:45.058Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a3/3f/888f8664d4f8e23a1363a5f50024be5216e199ab7ad0ba20988c7ed6d729/mlx-0.31.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:1b3fb0dda955b0d552ce57bdd6f42b3309ab21b067e40587d6848443d307e91f", size = 584796, upload-time = "2026-04-22T03:14:47.215Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dd/14/e9cd18b51f9e1dbcb060eec0fafc2d2428c8e1eacd9b0a02d7c5ce75b661/mlx-0.31.2-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:34b0171cd9eb5c43fdd82091f6135d6ccc5a065363a4a3e68fac64fb4e53d37c", size = 584790, upload-time = "2026-04-22T03:14:48.519Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ca/20/c6c5fb998c7834d094b2bfb9f003b5246cb270f0266da055c55546c34999/mlx-0.31.2-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:c05981684279a8935d58b0dde3ea5b02d210c3bad3319aa0e9934ec2df165752", size = 584795, upload-time = "2026-04-22T03:14:49.904Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0b/19/aca251d4c5f3532ce9c2c1e95ad76740d9c6c298f406f62d992f465b9be0/mlx-0.31.2-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:cd1f4189e5f1bc68735f44eb63ce98ae09d66ac75d7ab5b15a41afae7e9f0513", size = 627843, upload-time = "2026-04-22T03:14:51.351Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3e/2b/b89364883b98f21c2fe29e52d4ac8bc2fa2fe0d79293b36ec421efc1854a/mlx-0.31.2-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:53c8d57ffa9ce77f8355663be05014c0dd37280e57f19126fb0a24389a30684b", size = 685064, upload-time = "2026-04-22T03:14:52.75Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8b/f5/e63f6a9316ded2d14a8ebc7a9ca25734c784e8c54d064a78b4dceeacec0e/mlx-0.31.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:a13c9ce23c3deef6aa5a09315e7953e1a5dc311e851fa16fc74c81fb2509c0b9", size = 588417, upload-time = "2026-04-22T03:14:54.094Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/31/50/9d0c03ea3134cd85c132df7b0e4b75e6344bd8b4881a0b9c465cfa27f724/mlx-0.31.2-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:b0764bf11fc3a71dee988e19275eef67775cab63112d8bb7ef173ca8b2a1247c", size = 588421, upload-time = "2026-04-22T03:14:55.898Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ef/5b/d364cc793bcb504621313acb55627cf0d5403ab2e0a594aa081cdbe4591f/mlx-0.31.2-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:59ccbd0f0044d4f97f11ebcbf0c480bc9e962935fd96275f120954afea65be8a", size = 588384, upload-time = "2026-04-22T03:14:57.439Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1b/a4/e822202dd2e4e7d08671f2ecf7b6500af74f5bad5ceb27086b1aa6902f3a/mlx-0.31.2-cp314-cp314-manylinux_2_35_aarch64.whl", hash = "sha256:e81798c610f95a09c642c89214ba5c23b72ce18ce4728184aceabe7eddca33d7", size = 630473, upload-time = "2026-04-22T03:14:58.985Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/40/6f/da48d2d7a76e644d35438ef6f33c68755fdd382e2c546fd1804ccba01d04/mlx-0.31.2-cp314-cp314-manylinux_2_35_x86_64.whl", hash = "sha256:69fbc94bf53607a75af9eb3e22c354738a6fe4e25aa4e2b20934b009a4bba1f3", size = 685459, upload-time = "2026-04-22T03:15:00.45Z" }, +] + +[[package]] +name = "mlx-lm" +version = "0.31.3" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "jinja2" }, + { name = "mlx", marker = "sys_platform == 'darwin'" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "pyyaml" }, + { name = "sentencepiece" }, + { name = "transformers" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.31.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3f/69/fe3b783ebe999f3118234e1e940feb622518bfb1dea6ac5d13b1d36a8449/mlx_metal-0.31.2-py3-none-macosx_14_0_arm64.whl", hash = "sha256:b25385bcee18fc194092255b8b53b9a3d8489eb650e59160f1b57aadd07aa2dc", size = 40055588, upload-time = "2026-04-22T03:14:14.43Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4f/5d/4c690d5b93c30ba002656c37363159d978705bf8eb801b8481840fb942c2/mlx_metal-0.31.2-py3-none-macosx_15_0_arm64.whl", hash = "sha256:e9d4e5fce6ca10a87a0e388597f99519ad594d09e674708b5312bd8bd4f5997d", size = 40053220, upload-time = "2026-04-22T03:14:18.048Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/99/82/11fd62a8d7a3e96e5c43220b17de0151e3f10101f8bb3b865f5bd9cdd074/mlx_metal-0.31.2-py3-none-macosx_26_0_arm64.whl", hash = "sha256:84ffb60ee503f03eb684f5fb168d5cff31e2a16b7f27c1731eaf7662bd6e9b46", size = 55792151, upload-time = "2026-04-22T03:14:22.059Z" }, +] + +[[package]] +name = "nanobind" +version = "2.12.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c6/5c/3b69bc3933ad3c3668ba029ad410ba8ecfdc8ee7262ff1009f3304f3c562/nanobind-2.12.0.tar.gz", hash = "sha256:0ae77c1a88f27153fa57045ee00f7b0a7b06b1cd3df942e95a34b38c5d0a5bee", size = 1002704, upload-time = "2026-02-25T09:41:54.691Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/bf/1a54e3573736f3ad15fc599c5dde007937234652a1a7efd62573b4ce3a7e/nanobind-2.12.0-py3-none-any.whl", hash = "sha256:a10d3d88e691dcdf22696f9acd893fda3c5a05635763aea238829d274fcad480", size = 249512, upload-time = "2026-02-25T09:41:52.908Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.6" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/b3/49/ec46835a70be8fa6446c495126ac84fdb28cb2558e1620ffb87a10c8b64c/numpy-2.4.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0280e0356c0829a18d9de1cb7eee50ec22ca639878d7240307ca0943d73cd2c4", size = 16969194, upload-time = "2026-05-18T23:33:13.503Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/63/13/f9a8046535cb21deae82f8d03de9617e08882d274fad2539630761888228/numpy-2.4.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:8155154c7c691289fe18f510b5d4657c68c67989f293f0535a91360392ff6538", size = 6798936, upload-time = "2026-05-18T23:33:22.987Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/15/a7/9bc1cd626d7bf6869bfedf27b91b6ab5dd607758bf8e959d6fa80c6a59cb/numpy-2.4.6-cp311-cp311-win32.whl", hash = "sha256:ddea102b48f9e339f3948bf22040944184627a30fdf7f858667673b9c5f033c8", size = 6233945, upload-time = "2026-05-18T23:33:41.331Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c5/31/7fc6239c12bce7e931463251cca4426c465e1876ba3cc785402ef4dd8f4e/numpy-2.4.6-cp311-cp311-win_amd64.whl", hash = "sha256:1e254a00cdf42b1e4d5b3d68d33af63268d41340d8885df2ab6470f2e1500147", size = 12608406, upload-time = "2026-05-18T23:33:44.131Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/27/83/140f85a466595a16382996a1bf06b2b54bcd597488921b0c9daaeeda72af/numpy-2.4.6-cp311-cp311-win_arm64.whl", hash = "sha256:ed9749eef4cbd126da3dc1d6bcb3a57f5eb7ac6a6484146bdbf743f552dfc577", size = 10479528, upload-time = "2026-05-18T23:33:50.725Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/95/2a/3d7b5ac8aac24feaf9ad7ed58f45b0bbc06d37e4338ae84c9f2298b570f9/numpy-2.4.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:001fbb8e08d942dd57599e781f2472269ee7f2755fae407b4f67b2f0b17da3f1", size = 16689119, upload-time = "2026-05-18T23:33:54.065Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/f3/d4/9770d14ba719432bb90a421bfd443872ed0f70f7264b64bec12ea363d5fd/numpy-2.4.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:357cc07a6d7b0b182ff02249616a03742827ebb1277546b5c7cd7f7620a45698", size = 6551240, upload-time = "2026-05-18T23:34:02.852Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/82/bdab26d7438c6791ca31b7c024ca37c1eab8b726ba236129005cd4a06e45/numpy-2.4.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:511dbaf848decaaaf4b4ca48032619fb3138710c4bf7da7617765edad1ef96b0", size = 16684648, upload-time = "2026-05-18T23:34:29.41Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/ba/8c/ebd2a8f8a83541f8d38cc5667e8c2b69cecfd30da6e45693e8158857d44b/numpy-2.4.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:6180d8b35af935aed8ece3a85e0a43f87393ae0ac87c8d2c8bd2c993f7270ef3", size = 6546944, upload-time = "2026-05-18T23:34:38.484Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/f7/da/2ccc6c2fe8898dee01d90c75c5f5f914a23daf99e3e0f59516a08760c8b5/numpy-2.4.6-cp313-cp313-win32.whl", hash = "sha256:56b39e5e0622a09a25bf5baf62f4bcf0cb8a41ae6e2819cf49bbc5a74c083f91", size = 5961134, upload-time = "2026-05-18T23:34:55.618Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b5/cd/9cc4dc876fb065d5c220aae4d5e14826b2715331bb7618ce1fb07a679d99/numpy-2.4.6-cp313-cp313-win_amd64.whl", hash = "sha256:c4fc99836233ea196540b17ab0983aff60ed07941751930f5f4d05bc3b3b7359", size = 12318598, upload-time = "2026-05-18T23:34:58.928Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/39/1e/c0bcba1f8694116485fe28fd1be698c278fcda4141c5b0e53a2aed8b12a8/numpy-2.4.6-cp313-cp313-win_arm64.whl", hash = "sha256:a7c711e21628b52034bb5ab8d1bce291f752fcc5e92accc615778acee1ff4778", size = 10222272, upload-time = "2026-05-18T23:35:02.167Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/af/57/3917ab0fd97f271a8694513581b8a36c655f111c446852c302f04ccdb6fc/numpy-2.4.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:7265a2f3d436e54ef9f2b52b5c937e6be778781bd97a590319d7348f1c1ca997", size = 6646763, upload-time = "2026-05-18T23:35:11.459Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/8a/90/0ac3bc947217e66dec77e7cbc6a1979d1af70b6461b82f620d3bccd5e4c8/numpy-2.4.6-cp313-cp313t-win32.whl", hash = "sha256:29a287e0cf63ff528da061de6b9f64a4618da591ca1046aafc54062e40ca7eab", size = 6084971, upload-time = "2026-05-18T23:35:29.387Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/77/71/5673e351671a1d2bd6063b91b44f70c0affea7d1516fa7a6572941ba4aa1/numpy-2.4.6-cp313-cp313t-win_amd64.whl", hash = "sha256:25c692919ac5a01f170a3bfcd62d745b24fd095c353d50812637d6fcab442e75", size = 12458532, upload-time = "2026-05-18T23:35:32.175Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3f/88/19d3503c5046e688f049274b27a3ef3d771152fa80d3ba3d01a3dff61abe/numpy-2.4.6-cp313-cp313t-win_arm64.whl", hash = "sha256:1e978ec1e8bd0e0e4de6bb75de9d30cbb74db6b6a2bb727618613703ca0167dd", size = 10291881, upload-time = "2026-05-18T23:35:35.465Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f8/91/3ab2044d05fd16d343c5ac2e69b127f1b2854040dd20b193257c78028bd3/numpy-2.4.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:06ca2f61ec4385a07a6977c55ba998a4466c123642b4a32694d3128fce18c079", size = 16683458, upload-time = "2026-05-18T23:35:38.353Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8e/62/764ce66fa4147ae6d73071a3abf804ffe606f174618697c571acdf26a7c9/numpy-2.4.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:38efbc8de75c7a0fc1ac190162d892787f3f47b57cc291231aafee36b80982b7", size = 14704559, upload-time = "2026-05-18T23:35:42.14Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/61/23f27c172f022e04025b7dc2367f4d63c1a398120607ec896228649a6f48/numpy-2.4.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d581b735e177fdcdce6fed8e7e8880a3fb6ee4e3653a3ac6af01c6f4c03effc5", size = 5209716, upload-time = "2026-05-18T23:35:45.377Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/03/71/21cf70dc6ea3e3acb95fc53a265b2fc248b981f0194ceb5b475271b8809d/numpy-2.4.6-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:0a041d3d761dc3c35cc56ce0351506a02bcbc25f7b169f652435141a17db9096", size = 6543947, upload-time = "2026-05-18T23:35:47.926Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d5/91/64288395ee1799bd2e0b04a305dce9666da90c961e1f3fe982a05ee1c036/numpy-2.4.6-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40fdc1ae7125e518ea98e53e69a4ebc27e1fd50510c47b7ea130cf21e5e1d42b", size = 15685197, upload-time = "2026-05-18T23:35:50.863Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f3/eb/ebffaa97dc55502df69584a8f0dcf07f69a3e0b3e2323670a2722db9aa39/numpy-2.4.6-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a2c306dea656c12c68f51f4cea133cbe78ca7435eb28c735eac1d3ebe73be6e8", size = 16638245, upload-time = "2026-05-18T23:35:54.752Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b8/0b/54f9da33128d7e350fab89c7455902eeae70349ee52bddb448dc4a576f45/numpy-2.4.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:33111801a01c12a8a1e3721f0a9232f8cfc8ae2c6b7098167e6f623c6073f402", size = 17036587, upload-time = "2026-05-18T23:35:58.355Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b6/f0/fdebc1052db1cc37c64beb22072d67cd6d1c71adca1299f53dec2b5e20d3/numpy-2.4.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ae506e6902902557576a26ff33eda8695e7ecb3cb36c3b573a0765dee114ebdb", size = 18363226, upload-time = "2026-05-18T23:36:02.845Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/aa/b4/298628d98c72b57e57f7165ae6a481a1deaf6f3c28262a6e4c739c275930/numpy-2.4.6-cp314-cp314-win32.whl", hash = "sha256:aaf159caa35993cb1f56fb9b8e4610d35758e7ca005412eb1daa856a78c9c4b1", size = 6010196, upload-time = "2026-05-18T23:36:05.92Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/df/ac/46de6dda46478f7942f839e094970be2d4a861e005c4b3bf07c92e291a09/numpy-2.4.6-cp314-cp314-win_amd64.whl", hash = "sha256:b507f5c4c1d508876d1819b6bf9a49d365b96320b5d4993426b33a23ca4b8261", size = 12450334, upload-time = "2026-05-18T23:36:09.107Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/78/92/b8b798ac784102c0da830d2257d59358e3d3d90d1e2b3f2575dad976c5cf/numpy-2.4.6-cp314-cp314-win_arm64.whl", hash = "sha256:6f41ae150c4e32db4f3310cdaf64b1593a03dbabe29eec77fc9b50fe64061df6", size = 10495678, upload-time = "2026-05-18T23:36:12.766Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/30/34/ec28d1aa8115971537c01469ab2011ee96827930f0a124de1000cc2a7ed7/numpy-2.4.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ece3d2cfe132e7d51f44a832b303895e6f2d499c5e74dfbdb06ee246147a304a", size = 14823672, upload-time = "2026-05-18T23:36:16.473Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/16/bd/f6d1fede4e54e8042a7ff97bb495510f3c220f94bcd9e8b228e87c92cc0d/numpy-2.4.6-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:e3e5193ef5a3dc73bceee50f7fdc2c90dbb76c42df8d8fae3d1067a583df579e", size = 5328731, upload-time = "2026-05-18T23:36:19.767Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f4/f0/e105b9e2fd728a9910103884decd6951d9dd73896b914a98d9a231de02ee/numpy-2.4.6-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:17f9ade344e7d9b464a084d69bcf18fc691cb1db67c62ed80820bf4926d78f0e", size = 6649805, upload-time = "2026-05-18T23:36:22.266Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/82/dd/1206a7ca6ab15e3f02069707ca96222e202af681bb73756da7527f3cb837/numpy-2.4.6-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9cd5ffd25db4e7ba6a375693b3fc0fc1791ec636c17db3720da19bde7180ec43", size = 15730496, upload-time = "2026-05-18T23:36:25.713Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/51/e7/38d3ea825dcab85a591734decb2f6c67caa7c8367d374df1a1c3842f9b07/numpy-2.4.6-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7d92c3819208a60205a12a245c91ad70cb0a85336659b19b834205573ac8456e", size = 16679616, upload-time = "2026-05-18T23:36:29.652Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/93/b7/caabfdf53edf663e0b4eb74d7d405d83baef09eb5e83bcd32d601d72b93e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e85b752a1e912b70eaad4fafbd4d1238007ab221de2009b9a2f5ae7461239895", size = 17085145, upload-time = "2026-05-18T23:36:33.449Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f9/45/68d7c33a6bcf3e5aa3bdbd57a367e6f615286dfd6482f97e8ffeb734306e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:29cb7f67d10b479ff07c17d33e39f78c07f71c40ef30d63c153d340e96cd3fb4", size = 18403813, upload-time = "2026-05-18T23:36:37.369Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9c/50/0753655aa844c99cd9e018aacf76f130f1bd81d881bb74bc0aef5d73a8ba/numpy-2.4.6-cp314-cp314t-win32.whl", hash = "sha256:260a5d70215b61ab4fadf5c7baacd64821842975eea312125ed3c39a6391b063", size = 6156982, upload-time = "2026-05-18T23:36:40.817Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b2/d4/7c67becf668f973cb490cec3e98dfd799d866f9c989a54d355672cfa0db6/numpy-2.4.6-cp314-cp314t-win_amd64.whl", hash = "sha256:81a1cca95ed5bb92aa8b10dd2cdc9a0d3853a50fad926c28b5d7e8ea54389627", size = 12638908, upload-time = "2026-05-18T23:36:43.996Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/43/bb/e1c71a4295b1b1d1393d50dbb4f2a36283c6859d9d3892e84f00ec5a91d5/numpy-2.4.6-cp314-cp314t-win_arm64.whl", hash = "sha256:0c9136e14ed34a9e343a31c533d78a9813a69a3148332bce5e9821cb2f996e66", size = 10565867, upload-time = "2026-05-18T23:36:47.114Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/de/12/b422cc84439adc0d00de605bf4a308890ae5c26f2c71fbd73e5d08fbb0dd/numpy-2.4.6-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:55cced7c52e981362f708ad635198e97a752dfba412cc03c23bbf3bd8d5cd662", size = 16847511, upload-time = "2026-05-18T23:36:50.673Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/6a/ef/f66cc724fcc36c1e364c67f51ae9146090b8b584f27d58b97fdae3edd737/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_x86_64.whl", hash = "sha256:948424b06129ce883307e8cff868c31396d8dc7630a59c61d70d98dbe70f222c", size = 6708728, upload-time = "2026-05-18T23:36:59.575Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/15/ce/e5ec180bc41812edcd8daeb8639d205622c0e8c02259d8ab25a0201b3c2a/numpy-2.4.6-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:2803abfebfc990042cd494d8ce2d5f82e9d847af6d35ec486923aa19dbad5e73", size = 12504263, upload-time = "2026-05-18T23:37:09.715Z" }, +] + +[[package]] +name = "packaging" +version = "26.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.10.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/47/e4501f49c178ae1d9f4a75073fda4204f52647993f075a9db4d14930e0c5/platformdirs-4.10.0.tar.gz", hash = "sha256:31e761a6a0ca04faf7353ea759bdba55652be214725111e5aac52dfa29d4bef7", size = 31224, upload-time = "2026-05-28T03:32:53.587Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/81/e6/cd9575ac904136b3cbf7aa7ee819ef86eedb7274e46f230e94ea4342e729/platformdirs-4.10.0-py3-none-any.whl", hash = "sha256:fb516cdb12eb0d857d0cd85a7c57cea4d060bee4578d6cf5a14dfdf8cbf8784a", size = 22743, upload-time = "2026-05-28T03:32:52.175Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.35.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/fd/5b1491d9e4b586d621c54f4c36b888714164b6875f8d6afa3f9072906a51/protobuf-7.35.0.tar.gz", hash = "sha256:a2efd84605f41e559f1881b0912b44099d0a2ac9bf46b3474823f10fb393b0e6", size = 458677, upload-time = "2026-05-19T23:02:29.197Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/83/ee/93d06e358a4aa32280b00e722d3ea0a1f25fc3cc5778d80581c9cca2c10e/protobuf-7.35.0-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:66be6c513931c794fa92c080ffee41671390da3d79da219cf9c0c0907f035dda", size = 433225, upload-time = "2026-05-19T23:02:19.884Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8b/39/1c76c2da93f3c507e958e0aecee2391cc44d4625de6c728bbc555195b5a8/protobuf-7.35.0-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:fcbe42a4ac09d3ec9c987ddfcd956afd0b15f1ff613bd8371bde9405ffd5c8e5", size = 328847, upload-time = "2026-05-19T23:02:22.3Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/91/1a/39f7ce90a238c1a987a4d81ec26379e02ca0aff367de68e4a1fa474215b9/protobuf-7.35.0-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:4cbf5cc286130e06a6c9bbefac442431173906dfcc979712183d4adcc01b37ee", size = 344030, upload-time = "2026-05-19T23:02:23.591Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/70/5b/6baf9008817964454055ff3fe65f1de0b5f1e26c80c82f7fb108b7cd4ea3/protobuf-7.35.0-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:6c0f98f10c8a05ea30f8993dfef2de093d27b490fdae78bb60c8343795d55011", size = 327130, upload-time = "2026-05-19T23:02:24.637Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8e/e5/e46adb0badc388bfb84877a5f9f026aff63f60e611016cf64dbe77e05446/protobuf-7.35.0-cp310-abi3-win32.whl", hash = "sha256:4c4617b83ade0e279d1d2bfe04025a1adb87f9ed657de038620dc0ff959357f6", size = 428946, upload-time = "2026-05-19T23:02:25.741Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a7/ab/547fbd9e16d879dd13c167478f8ae0a83a428008ca07a5e06acdc23ad473/protobuf-7.35.0-cp310-abi3-win_amd64.whl", hash = "sha256:f05bcadf9a2a6b8dda047007075135fb7d08c73d9177aabc067e1be46881a201", size = 439996, upload-time = "2026-05-19T23:02:26.808Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b8/ef/50433d346c56657a70d27f156c7b349ac59a068b01de4eb796e747eecc43/protobuf-7.35.0-py3-none-any.whl", hash = "sha256:c13f325cf242bad135c350629eeb5d54b24228eb472fb3e2e9ebbd4c5dc20ca0", size = 171659, upload-time = "2026-05-19T23:02:27.842Z" }, +] + +[[package]] +name = "pydantic" +version = "2.13.4" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.4" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5c/fa/6d7708d2cfc1a832acb6aeb0cd16e801902df8a0f583bb3b4b527fde022e/pydantic_core-2.46.4-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594", size = 2111872, upload-time = "2026-05-06T13:40:27.596Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ae/6f/aa064a3e74b5745afbdf250594f38e7ead05e2d651bcb35994b9417a0d4d/pydantic_core-2.46.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c", size = 1948255, upload-time = "2026-05-06T13:39:12.574Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/43/3a/41114a9f7569b84b4d84e7a018c57c56347dac30c0d4a872946ec4e36c46/pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826", size = 1972827, upload-time = "2026-05-06T13:38:19.841Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ef/25/1ab42e8048fe551934d9884e8d64daa7e990ad386f310a15981aeb6a5b08/pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04", size = 2041051, upload-time = "2026-05-06T13:38:10.447Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/94/c2/1a934597ddf08da410385b3b7aae91956a5a76c635effef456074fad7e88/pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e", size = 2221314, upload-time = "2026-05-06T13:40:13.089Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/02/6d/9e8ad178c9c4df27ad3c8f25d1fe2a7ab0d2ba0559fad4aee5d3d1f16771/pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3", size = 2285146, upload-time = "2026-05-06T13:38:59.224Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/80/50/540cd3aeefc041beb111125c4bff779831a2111fc6b15a9138cda277d32c/pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4", size = 2089685, upload-time = "2026-05-06T13:38:17.762Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6b/a4/b440ad35f05f6a38f89fa0f149accb3f0e02be94ca5e15f3c449a61b4bc9/pydantic_core-2.46.4-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398", size = 2115420, upload-time = "2026-05-06T13:37:58.195Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/99/61/de4f55db8dfd57bfdfa9a12ec90fe1b57c4f41062f7ca86f08586b3e0ac0/pydantic_core-2.46.4-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3", size = 2165122, upload-time = "2026-05-06T13:37:01.167Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f7/52/7c529d7bdb2d1068bd52f51fe32572c8301f9a4febf1948f10639f1436f5/pydantic_core-2.46.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848", size = 2182573, upload-time = "2026-05-06T13:38:45.04Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/37/b3/7c40325848ba78247f2812dcf9c7274e38cd801820ca6dd9fe63bcfb0eb4/pydantic_core-2.46.4-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3", size = 2317139, upload-time = "2026-05-06T13:37:15.539Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d9/37/f913f81a657c865b75da6c0dbed79876073c2a43b5bd9edbe8da785e4d49/pydantic_core-2.46.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109", size = 2360433, upload-time = "2026-05-06T13:37:30.099Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c4/67/6acaa1be2567f9256b056d8477158cac7240813956ce86e49deae8e173b4/pydantic_core-2.46.4-cp311-cp311-win32.whl", hash = "sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda", size = 1985513, upload-time = "2026-05-06T13:38:15.669Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/aa/e6/c505f83dfeda9a2e5c995cfd872949e4d05e12f7feb3dca72f633daefa94/pydantic_core-2.46.4-cp311-cp311-win_amd64.whl", hash = "sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33", size = 2071114, upload-time = "2026-05-06T13:40:35.416Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0f/da/7a263a96d965d9d0df5e8de8a475f33495451117035b09acb110288c381f/pydantic_core-2.46.4-cp311-cp311-win_arm64.whl", hash = "sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d", size = 2044298, upload-time = "2026-05-06T13:38:29.754Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ce/8c/af022f0af448d7747c5154288d46b5f2bc5f17366eaa0e23e9aa04d59f3b/pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2", size = 2106158, upload-time = "2026-05-06T13:38:57.215Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/95/6195171e385007300f0f5574592e467c568becce2d937a0b6804f218bc49/pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f", size = 1951724, upload-time = "2026-05-06T13:37:02.697Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8e/bc/f47d1ff9cbb1620e1b5b697eef06010035735f07820180e74178226b27b3/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7", size = 1975742, upload-time = "2026-05-06T13:37:09.448Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5b/11/9b9a5b0306345664a2da6410877af6e8082481b5884b3ddd78d47c6013ce/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7", size = 2052418, upload-time = "2026-05-06T13:37:38.234Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f1/b7/a65fec226f5d78fc39f4a13c4cc0c768c22b113438f60c14adc9d2865038/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712", size = 2232274, upload-time = "2026-05-06T13:38:27.753Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/68/f0/92039db98b907ef49269a8271f67db9cb78ae2fc68062ef7e4e77adb5f61/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4", size = 2309940, upload-time = "2026-05-06T13:38:05.353Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/22/37/a8aca44d40d737dde2bc05b3c6c07dff0de07ce6f82e9f3167aeaf4d5dea/pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987", size = 2136854, upload-time = "2026-05-06T13:40:22.59Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/99/fcef1b79238c06a8cbec70819ac722ba76e02bc8ada9b0fd66eba40da01b/pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b", size = 2180306, upload-time = "2026-05-06T13:40:10.666Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ae/6c/fc44000918855b42779d007ae63b0532794739027b2f417321cddbc44f6a/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458", size = 2190044, upload-time = "2026-05-06T13:40:43.231Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6b/65/d9cadc9f1920d7a127ad2edba16c1db7916e59719285cd6c94600b0080ba/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b", size = 2329133, upload-time = "2026-05-06T13:39:57.365Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/47/bd/6f2fc8188f31bf10590f1e98e7b306336161fac930a8c514cd7bd828c7dc/pydantic_core-2.46.4-cp312-cp312-win32.whl", hash = "sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894", size = 1974823, upload-time = "2026-05-06T13:40:47.985Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/40/8c/985c1d41ea1107c2534abd9870e4ed5c8e7669b5c308297835c001e7a1c4/pydantic_core-2.46.4-cp312-cp312-win_amd64.whl", hash = "sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89", size = 2072919, upload-time = "2026-05-06T13:39:21.153Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c4/ba/f463d006e0c47373ca7ec5e1a261c59dc01ef4d62b2657af925fb0deee3a/pydantic_core-2.46.4-cp312-cp312-win_arm64.whl", hash = "sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a", size = 2027604, upload-time = "2026-05-06T13:39:03.753Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/51/a2/5d30b469c5267a17b39dec53208222f76a8d351dfac4af661888c5aee77d/pydantic_core-2.46.4-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008", size = 2106306, upload-time = "2026-05-06T13:37:48.029Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c1/81/4fa520eaffa8bd7d1525e644cd6d39e7d60b1592bc5b516693c7340b50f1/pydantic_core-2.46.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4", size = 1951906, upload-time = "2026-05-06T13:37:17.012Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/03/d5/fd02da45b659668b05923b17ba3a0100a0a3d5541e3bd8fcc4ecb711309e/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76", size = 1976802, upload-time = "2026-05-06T13:37:35.113Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/21/f2/95727e1368be3d3ed485eaab7adbd7dda408f33f7a36e8b48e0144002b91/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3", size = 2052446, upload-time = "2026-05-06T13:37:12.313Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9c/86/5d99feea3f77c7234b8718075b23db11532773c1a0dbd9b9490215dc2eeb/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76", size = 2232757, upload-time = "2026-05-06T13:39:01.149Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d2/3a/508ac615935ef7588cf6d9e9b91309fdc2da751af865e02a9098de88258c/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4", size = 2309275, upload-time = "2026-05-06T13:37:41.406Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/07/f8/41db9de19d7987d6b04715a02b3b40aea467000275d9d758ffaa31af7d50/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a", size = 2094467, upload-time = "2026-05-06T13:39:18.847Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2c/e2/f35033184cb11d0052daf4416e8e10a502ea2ac006fc4f459aee872727d1/pydantic_core-2.46.4-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262", size = 2134417, upload-time = "2026-05-06T13:40:17.944Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7e/7b/6ceeb1cc90e193862f444ebe373d8fdf613f0a82572dde03fb10734c6c71/pydantic_core-2.46.4-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e", size = 2179782, upload-time = "2026-05-06T13:40:32.618Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/f2/c8d7773ede6af08036423a00ae0ceffce266c3c52a096c435d68c896083f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd", size = 2188782, upload-time = "2026-05-06T13:36:51.018Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/59/31/0c864784e31f09f05cdd87606f08923b9c9e7f6e51dd27f20f62f975ce9f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be", size = 2328334, upload-time = "2026-05-06T13:40:37.764Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c2/eb/4f6c8a41efa30baa755590f4141abf3a8c370fab610915733e74134a7270/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d", size = 2372986, upload-time = "2026-05-06T13:39:34.152Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5b/24/b375a480d53113860c299764bfe9f349a3dc9108b3adc0d7f0d786492ebf/pydantic_core-2.46.4-cp313-cp313-win32.whl", hash = "sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb", size = 1973693, upload-time = "2026-05-06T13:37:55.072Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7e/e8/cff247591966f2d22ec8c003cd7587e27b7ba7b81ab2fb888e3ab75dc285/pydantic_core-2.46.4-cp313-cp313-win_amd64.whl", hash = "sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292", size = 2071819, upload-time = "2026-05-06T13:38:49.139Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c6/1a/f4aee670d5670e9e148e0c82c7db98d780be566c6e6a97ee8035528ca0b3/pydantic_core-2.46.4-cp313-cp313-win_arm64.whl", hash = "sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d", size = 2027411, upload-time = "2026-05-06T13:40:45.796Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8d/74/228a26ddad29c6672b805d9fd78e8d251cd04004fa7eed0e622096cd0250/pydantic_core-2.46.4-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb", size = 2102079, upload-time = "2026-05-06T13:38:41.019Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ad/1f/8970b150a4b4365623ae00fc88603491f763c627311ae8031e3111356d6e/pydantic_core-2.46.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462", size = 1952179, upload-time = "2026-05-06T13:36:59.812Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/95/30/5211a831ae054928054b2f79731661087a2bc5c01e825c672b3a4a8f1b3e/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9", size = 1978926, upload-time = "2026-05-06T13:37:39.933Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/57/e9/689668733b1eb67adeef047db3c2e8788fcf65a7fd9c9e2b46b7744fe245/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4", size = 2046785, upload-time = "2026-05-06T13:38:01.995Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/d9/6715260422ff50a2109878fd24d948a6c3446bb2664f34ee78cd972b3acd/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914", size = 2228733, upload-time = "2026-05-06T13:40:50.371Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/ae/fdb2f64316afca925640f8e70bb1a564b0ec2721c1389e25b8eb4bf9a299/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28", size = 2307534, upload-time = "2026-05-06T13:37:21.531Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/89/1d/8eff589b45bb8190a9d12c49cfad0f176a5cbd1534908a6b5125e2886239/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b", size = 2099732, upload-time = "2026-05-06T13:39:31.942Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/06/d5/ee5a3366637fee41dee51a1fc91562dcf12ddbc68fda34e6b253da2324bb/pydantic_core-2.46.4-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c", size = 2129627, upload-time = "2026-05-06T13:37:25.033Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/94/33/2414be571d2c6a6c4d08be21f9292b6d3fdb08949a97b6dfe985017821db/pydantic_core-2.46.4-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb", size = 2179141, upload-time = "2026-05-06T13:37:14.046Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7b/79/7daa95be995be0eecc4cf75064cb33f9bbbfe3fe0158caf2f0d4a996a5c7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898", size = 2184325, upload-time = "2026-05-06T13:36:53.615Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9f/cb/d0a382f5c0de8a222dc61c65348e0ce831b1f68e0a018450d31c2cace3a5/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e", size = 2323990, upload-time = "2026-05-06T13:40:29.971Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/db/d9ba624cc4a5aced1598e88c04fdbd8310c8a69b9d38b9a3d39ce3a61ed7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519", size = 2369978, upload-time = "2026-05-06T13:37:23.027Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f2/20/d15df15ba918c423461905802bfd2981c3af0bfa0e40d05e13edbfa48bc3/pydantic_core-2.46.4-cp314-cp314-win32.whl", hash = "sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4", size = 1966354, upload-time = "2026-05-06T13:38:03.499Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fc/b6/6b8de4c0a7d7ab3004c439c80c5c1e0a3e8d78bbae19379b01960383d9e5/pydantic_core-2.46.4-cp314-cp314-win_amd64.whl", hash = "sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac", size = 2072238, upload-time = "2026-05-06T13:39:40.807Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/32/36/51eb763beec1f4cf59b1db243a7dcc39cbb41230f050a09b9d69faaf0a48/pydantic_core-2.46.4-cp314-cp314-win_arm64.whl", hash = "sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a", size = 2018251, upload-time = "2026-05-06T13:37:26.72Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/91/855af51d625b23aa987116a19e231d2aaef9c4a415273ddc189b79a45fee/pydantic_core-2.46.4-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0", size = 2099593, upload-time = "2026-05-06T13:39:47.682Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/1b/8784a54c65edb5f49f0a14d6977cf1b209bba85a4c77445b255c2de58ab3/pydantic_core-2.46.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d", size = 1935226, upload-time = "2026-05-06T13:40:40.428Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/e7/1955d28d1afc56dd4b3ad7cc0cf39df1b9852964cf16e5d13912756d6d6b/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b", size = 1974605, upload-time = "2026-05-06T13:37:32.029Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/93/e2/3fedbf0ba7a22850e6e9fd78117f1c0f10f950182344d8a6c535d468fdd8/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000", size = 2030777, upload-time = "2026-05-06T13:38:55.239Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f8/61/46be275fcaaba0b4f5b9669dd852267ce1ff616592dccf7a7845588df091/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e", size = 2236641, upload-time = "2026-05-06T13:37:08.096Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/db/12e93e46a8bac9988be3c016860f83293daea8c716c029c9ace279036f2f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd", size = 2286404, upload-time = "2026-05-06T13:40:20.221Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e2/4a/4d8b19008f38d31c53b8219cfedc2e3d5de5fe99d90076b7e767de29274f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3", size = 2109219, upload-time = "2026-05-06T13:38:12.153Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/88/70/3cbc40978fefb7bb09c6708d40d4ad1a5d70fd7213c3d17f971de868ec1f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7", size = 2110594, upload-time = "2026-05-06T13:40:02.971Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9d/20/b8d36736216e29491125531685b2f9e61aa5b4b2599893f8268551da3338/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff", size = 2159542, upload-time = "2026-05-06T13:39:27.506Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1d/a2/367df868eb584dacf6bf82a389272406d7178e301c4ac82545ab98bc2dd9/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424", size = 2168146, upload-time = "2026-05-06T13:38:31.93Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c1/b8/4460f77f7e201893f649a29ab355dddd3beee8a97bcb1a320db414f9a06e/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6", size = 2306309, upload-time = "2026-05-06T13:37:44.717Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/64/c4/be2639293acd87dc8ddbcec41a73cee9b2ebf996fe6d892a1a74e88ad3f7/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565", size = 2369736, upload-time = "2026-05-06T13:37:05.645Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/30/a6/9f9f380dbb301f67023bf8f707aaa75daadf84f7152d95c410fd7e81d994/pydantic_core-2.46.4-cp314-cp314t-win32.whl", hash = "sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02", size = 1955575, upload-time = "2026-05-06T13:38:51.116Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/40/1f/f1eb9eb350e795d1af8586289746f5c5677d16043040d63710e22abc43c9/pydantic_core-2.46.4-cp314-cp314t-win_amd64.whl", hash = "sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5", size = 2051624, upload-time = "2026-05-06T13:38:21.672Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f6/d2/42dd53d0a85c27606f316d3aa5d2869c4e8470a5ed6dec30e4a1abe19192/pydantic_core-2.46.4-cp314-cp314t-win_arm64.whl", hash = "sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596", size = 2017325, upload-time = "2026-05-06T13:40:52.723Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ee/a4/73995fd4ebbb46ba0ee51e6fa049b8f02c40daebb762208feda8a6b7894d/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c", size = 2111589, upload-time = "2026-05-06T13:37:10.817Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/7f/f37d3a5e8bfcc2e403f5c57a730f2d815693fb42119e8ea48b3789335af1/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b", size = 1944552, upload-time = "2026-05-06T13:36:56.717Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/15/3c/d7eb777b3ff43e8433a4efb39a17aa8fd98a4ee8561a24a67ef5db07b2d6/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b", size = 1982984, upload-time = "2026-05-06T13:39:06.207Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/63/87/70b9f40170a81afd55ca26c9b2acb25c20d64bcfbf888fafecb3ba077d4c/pydantic_core-2.46.4-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea", size = 2138417, upload-time = "2026-05-06T13:39:45.476Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9d/1d/8987ad40f65ae1432753072f214fb5c74fe47ffbd0698bb9cbbb585664f8/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7", size = 2095527, upload-time = "2026-05-06T13:39:52.283Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/64/d3/84c282a7eee1d3ac4c0377546ef5a1ea436ce26840d9ac3b7ed54a377507/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df", size = 1936024, upload-time = "2026-05-06T13:40:15.671Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/ca/eac61596cdeb4d7e174d3dc0bd8a6238f14f75f97a24e7b7db4c7e7340a0/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526", size = 1990696, upload-time = "2026-05-06T13:38:34.717Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/11/cb/428de0385b6c8d44b716feba566abfacfbd23ee3c4439faa789a1456242f/pydantic_core-2.46.4-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0", size = 2112782, upload-time = "2026-05-06T13:37:04.016Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0b/b5/6a17bdadd0fc1f170adfd05a20d37c832f52b117b4d9131da1f41bb097ce/pydantic_core-2.46.4-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7", size = 1952146, upload-time = "2026-05-06T13:39:43.092Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2a/dc/03734d80e362cd43ef65428e9de77c730ce7f2f11c60d2b1e1b39f0fbf99/pydantic_core-2.46.4-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2", size = 2134492, upload-time = "2026-05-06T13:36:58.124Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/de/df/5e5ffc085ed07cc22d298134d3d911c63e91f6a0eb91fe646750a3209910/pydantic_core-2.46.4-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9", size = 2156604, upload-time = "2026-05-06T13:37:49.88Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/81/44/6e112a4253e56f5705467cbab7ab5e91ee7398ba3d56d358635958893d3e/pydantic_core-2.46.4-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf", size = 2183828, upload-time = "2026-05-06T13:37:43.053Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ac/ad/5565071e937d8e752842ac241463944c9eb14c87e2d269f2658a5bd05e98/pydantic_core-2.46.4-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30", size = 2310000, upload-time = "2026-05-06T13:37:56.694Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4f/c3/66883a5cec183e7fba4d024b4cbbe61851a63750ef606b0afecc46d1f2bf/pydantic_core-2.46.4-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc", size = 2361286, upload-time = "2026-05-06T13:40:05.667Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4b/2d/69abac8f838090bbecd5df894befb2c2619e7996a98ddb949db9f3b93225/pydantic_core-2.46.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983", size = 2193071, upload-time = "2026-05-06T13:38:08.682Z" }, +] + +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[package]] +name = "pytest" +version = "9.0.3" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" }, +] + +[[package]] +name = "pytest-asyncio" +version = "1.4.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "pytest" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/43/7c/d36d04db312ecf4298932ef77e6e4a9e8ad017906e24e34f0b0c361a2473/pytest_asyncio-1.4.0.tar.gz", hash = "sha256:c6c0d2259945122819f171a32ecea2c349ead889ee28176caaf492143424be42", size = 58514, upload-time = "2026-05-26T09:56:04.083Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/03/e2/08a497ef684b88559c9cc5f4ad53a37e7b99e727094a86d6ea32536d5d3c/pytest_asyncio-1.4.0-py3-none-any.whl", hash = "sha256:933ca923a23075a87fb7070c0ec272a6848489824d887c85c812670932835aa1", size = 16930, upload-time = "2026-05-26T09:56:02.576Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/6d/16/a95b6757765b7b031c9374925bb718d55e0a9ba8a1b6a12d25962ea44347/pyyaml-6.0.3-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:44edc647873928551a01e7a563d7452ccdebee747728c1080d881d68af7b997e", size = 185826, upload-time = "2025-09-25T21:31:58.655Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/45/91/47a6e1c42d9ee337c4839208f30d9f09caa9f720ec7582917b264defc875/pyyaml-6.0.3-cp311-cp311-win32.whl", hash = "sha256:8098f252adfa6c80ab48096053f512f2321f0b998f98150cea9bd23d83e1467b", size = 142543, upload-time = "2025-09-25T21:32:08.95Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/e3/ea007450a105ae919a72393cb06f122f288ef60bba2dc64b26e2646fa315/pyyaml-6.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:9f3bfb4965eb874431221a3ff3fdcddc7e74e3b07799e0e84ca4a0f867d449bf", size = 158763, upload-time = "2025-09-25T21:32:09.96Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, +] + +[[package]] +name = "regex" +version = "2026.5.9" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dc/0e/49aee608ad09480e7fd276898c99ec6192985fa331abe4eb3a986094490b/regex-2026.5.9.tar.gz", hash = "sha256:a8234aa23ec39894bfe4a3f1b85616a7032481964a13ac6fc9f10de4f6fca270", size = 416074, upload-time = "2026-05-09T23:15:19.37Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c2/dc/c1f2df4027e82fc54b5a473e4b250f5139faca49a0fbe29a48668d228f34/regex-2026.5.9-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ccf5249114cc3e772ecdd88a98a86eca0fd74c61ce32a94743758c083fc05d48", size = 489445, upload-time = "2026-05-09T23:12:06.111Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/03/d2/59f01110660081cce9c0bc30ebd0b5ee250dacf658e3248ed92f01e0e8ee/regex-2026.5.9-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:46f1326ca6e65b0879d23ca302c0f2415aad42ff0309b9c818e7949fe19a41d8", size = 291271, upload-time = "2026-05-09T23:12:07.731Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/58/b6/14b2c84ff90ddb370c81d27503f4a0fcf071496416f4855f6cc8c5d81c35/regex-2026.5.9-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ef31cbfe458e21c6122ba8150ff060e0c7789ed0d26eb423f25472584920b555", size = 289212, upload-time = "2026-05-09T23:12:09.266Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/03/d0/4db86529117320de0c84afd90e70bb47434625875e34fcef9d8c127c5b16/regex-2026.5.9-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:992604d02e6d9c6d786c24a706a71ecffe1020fc1ef264044474cd81fa2c3919", size = 792310, upload-time = "2026-05-09T23:12:11.416Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/07/78/fe4800cd322f862ecffd2d553409b20d80650e5ed71b9d178f853d020b82/regex-2026.5.9-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c9411dd64ca95477225734a93dfc8583b51916b8d5942f99d6cac21e09965451", size = 861721, upload-time = "2026-05-09T23:12:13.681Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b5/d0/b3618a895dd8feb897c61bb2954edd265e1767d82a01d53065d5871127a3/regex-2026.5.9-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3dd4a3ff360dfb836fecdb93a4598f9d6e2ac81e3e397125145c6221bf58cf4c", size = 906460, upload-time = "2026-05-09T23:12:15.443Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/33/6f/1481597e859ef19508b345eec4afd1416ed6e6b459c75a64026ef193aecf/regex-2026.5.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2a661a7d270a61f7cf460caee8b9fa2d5ef9e5c681234bcb9e0fe14f488e7dfc", size = 799843, upload-time = "2026-05-09T23:12:16.892Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/59/955734c803f59108deccba3597ae440c76b62a652733c0006e6243758420/regex-2026.5.9-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f079e50a0d3cc3cd5091fa9ff45869a2e6b2cd35895731edafb0327901a8d86d", size = 773610, upload-time = "2026-05-09T23:12:19.127Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/68/8f/70c04a236d651c81881dac42ef8538bddda6121434509d0a22d9e601503b/regex-2026.5.9-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4ebe8f0b5ec5a5024dc4a4c59f444c4e9afc5f2abdbb8962065b75d27fb971f9", size = 781645, upload-time = "2026-05-09T23:12:20.806Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1d/96/05c7434d88185e5d27fe54aeb74df86bd77cd79f52f0b4eae54faa8fea70/regex-2026.5.9-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:97cf3bc1b7d7d2306772ec07366c80d9df00ff79e79cea32898883a646d2fae2", size = 854473, upload-time = "2026-05-09T23:12:22.465Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4e/c1/6e3d8202d981f3117004bf341ee74893ba4ba8a9fbaf4b94615846550a08/regex-2026.5.9-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:0f9eede6a5cbdc02d4978090186390936e1776a7d1359b21e41014c609880bcf", size = 763311, upload-time = "2026-05-09T23:12:24.351Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/93/c7/e7737f1526b3fb32bd4c337fd6c71c3ebb5c8296fc34d11197e0955d2e35/regex-2026.5.9-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:01f0f5f55f4b64dacec85dc116d3c05fd23ad3ff037bbc73a2085775953c2611", size = 844593, upload-time = "2026-05-09T23:12:26.341Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a5/27/0daffb1a535bb39f422c3d200f4ab023c71110ad66a32b366bee708baba0/regex-2026.5.9-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1268eddd8486dc561d08eee1156e40aa3a8fe10f4bdec8fa653b455fcbffd12c", size = 789167, upload-time = "2026-05-09T23:12:27.975Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ce/fc/294fe4fac4f2ed67207b17471815870c1c45b3a489e08e0ac96daea16ef6/regex-2026.5.9-cp311-cp311-win32.whl", hash = "sha256:8676474c07469d6f33dd1085ca2cd45f65785f32518f2b20e36d9953ca07f994", size = 266249, upload-time = "2026-05-09T23:12:30.141Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d0/b0/8dce459f6245bcf8f6e9f23ac9569f1a0f15c131cc0745e82b43226204cf/regex-2026.5.9-cp311-cp311-win_amd64.whl", hash = "sha256:246de9d60aa3f8538b519834dd95cbf276ea263d6a7bd5a3666dc3fa0230505b", size = 278423, upload-time = "2026-05-09T23:12:31.676Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/db/8d/f9aeff6ad63a3ef720386f2907e6d34a35a510a6e498ebad28b0fb3f6ab6/regex-2026.5.9-cp311-cp311-win_arm64.whl", hash = "sha256:d726ca3f0d76969bf1e8e477d160d3d666bbf999f6860bd314889e5345782046", size = 270420, upload-time = "2026-05-09T23:12:33.194Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/50/9b/6550044bc44e17c84d312c031c2ec42fbdb6a4ec4e29093be3a172d08772/regex-2026.5.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57eeeb05db7979413dec5438f2db21d7ecbba787cde7a711df1a6f6df672aa06", size = 490451, upload-time = "2026-05-09T23:12:34.72Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/95/fc7ba4303b5a0f92446a12ee6778ef2c6c799233f5060042a31bf390cfe9/regex-2026.5.9-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:398c521292f4c7fb807001dcd54694d3a1fcafc179a36ad9cc56f98df85930b6", size = 292112, upload-time = "2026-05-09T23:12:36.285Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/54/4b/ee27938d1b2c443e89a9a10e00d2d19aa5ee300cd3d61140644e93bb083e/regex-2026.5.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f7a7c26137296beba7784de6eba69c6a93a63ccebc385e4962fe67e267a91225", size = 289599, upload-time = "2026-05-09T23:12:38.089Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d8/dd/ba103dc19614e25f3880800ca67ce093d6e21b325d72b8383c7bf906e9fa/regex-2026.5.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6441cc660d76107934a09c22167200839a0e89604a6297f78a974e66e931d2c0", size = 796732, upload-time = "2026-05-09T23:12:40.062Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cf/e7/f035b4fd858b050b0080bf302968dc0f59ba34e391872d54936758e6844e/regex-2026.5.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:91328f1c23d47595ca3ef0a7557fa129c5a23404b775c770697d2f35b33e0107", size = 865440, upload-time = "2026-05-09T23:12:42.059Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0a/51/8cd301ecc899aea28124357f729f4272f44de7806fc7ca02490bfbe253e8/regex-2026.5.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:93a7860539414dddaefba2b40f8771765ae17949d4c7182b876ce429e11a8309", size = 912329, upload-time = "2026-05-09T23:12:44.373Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cc/1e/3fbe2fa1e8cebd62f3bb7d3321cff1640aca2e240b51d9bd624aad949260/regex-2026.5.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dd2810d22146b6d838acc5ec15602cb6b47920aa4e33015df3868eedfd20bab8", size = 801239, upload-time = "2026-05-09T23:12:46.268Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/17/2f/6f6008682bf2cf98040a0d3153a8e557b6ab728d7713d045cee4ce544ab8/regex-2026.5.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:daff2bdbaf1d23e52fdff7c0b7bc2048b68f978df6a4d107ac981f94caef2e66", size = 777054, upload-time = "2026-05-09T23:12:48.051Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/2b/eee0d20a6842ba04df4b8847a920b57ef56853f14ef85405473e586b605a/regex-2026.5.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4eeb011098fcb77af513dcef521a3dbecbf8849b1e38940759d293b7a93f5026", size = 785098, upload-time = "2026-05-09T23:12:49.851Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4a/98/6fc1e6410feefb92159edaed5041992bfe390e8d26c721865434acbca558/regex-2026.5.9-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:ea9c8ecfa1b73c73b626534d6626e5340d429630943672b8480724f44e84b962", size = 860095, upload-time = "2026-05-09T23:12:51.666Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/a3/bd855e0f2cb1a978ecf6fa6bb69632dd9c3f6ea3b81cde62fde14c9daec7/regex-2026.5.9-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:cd2846168eb9ee3c513902bc8225409cb1caab31d04728b145171fa1625d9621", size = 765762, upload-time = "2026-05-09T23:12:53.413Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dc/66/0ae8c092e60b14c79d24f8e0b7f0aea5bfbffdcab00b5483d13404d3c3a5/regex-2026.5.9-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:39617fb0cde9c0e6306dc70e3bfc096f3da793219879f7ae7aa341a69fbdcf6d", size = 852100, upload-time = "2026-05-09T23:12:55.256Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/21/de/8dfde60fc1b21c946a893ba273403b72617edb261370cb1087099a83f088/regex-2026.5.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fd03c4f0e33280d15cae17159b899245d6b7c53d21def19b263b39655061f5ce", size = 789479, upload-time = "2026-05-09T23:12:57.573Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c3/1c/bdcc98f9a4af4fdd166c74941174619ccff4726d3ce32faa8e9a2ecd38dd/regex-2026.5.9-cp312-cp312-win32.whl", hash = "sha256:164eba9b755ea6f244b0d881196fbc1fac09714e9782c9e2732b813142033c8e", size = 266699, upload-time = "2026-05-09T23:12:59.14Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/78/87/240d36864f9e48ace85f72e79ced97ceb7f27ce87739a947dcb834b4e6bc/regex-2026.5.9-cp312-cp312-win_amd64.whl", hash = "sha256:86f40a5d6444db30a125c9c9177e6b25dad981cbc37451fd838f145e6edac92e", size = 277783, upload-time = "2026-05-09T23:13:00.789Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4f/b5/7b30f312b0669dff5beebe5b0989dc2d1a312b1a44fab852199c387a5b96/regex-2026.5.9-cp312-cp312-win_arm64.whl", hash = "sha256:96f5f58b54a063d7ea9dca08e1cf57bfe10499c4d579ee672da284f57f5f0070", size = 270513, upload-time = "2026-05-09T23:13:02.426Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/aa/da/797e91ecec6f84135da778ddce78c20e0af5d2a15c26f87a81bc3eadb6db/regex-2026.5.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:d626b84406444b165fc0ba981604edea39f0588ff1f92baa23fe50799ea9afdb", size = 490303, upload-time = "2026-05-09T23:13:04.382Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/44/da/bf30abaaa737b58f4a4b8c4a03659e02fd92092c822e0197ed9e0daab917/regex-2026.5.9-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d7bdc0ab8f3dd7e1b4f9ab88634e13374669db86bb3c72e8292f07ae313f539f", size = 292019, upload-time = "2026-05-09T23:13:06.022Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2d/e7/d0eaf5713828417b9e5648cf81fa9bacd4961f6ab98c380c2034f8716e35/regex-2026.5.9-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a8820737949116ffff55fe18f9fc644530063ba6ebfcb8314239416e78f1347c", size = 289468, upload-time = "2026-05-09T23:13:08.214Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d3/9b/b3fdd62b003baa1a9b593cd8c8699c9651c2e80cc21a5c715707983c42d7/regex-2026.5.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0fbdbac82cb3e4450d0ccde7d7a35607f4cb2dd9fba4b8b69bfaf8c9fa6aed", size = 796749, upload-time = "2026-05-09T23:13:10.573Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d4/30/66ab84588765f5b4b271a9ca09ef7ce2b87caa95176ec3d2ad65d7bc4902/regex-2026.5.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:57e8915c7986aa33d25e4d3629cef711cd2863f2961b10409f0c04cb8b7d9020", size = 865445, upload-time = "2026-05-09T23:13:12.523Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1a/89/f05169e8588aac365f35ffc7f3bc3184f095ef4cfded7cfaa3c7fd5dbd89/regex-2026.5.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:508f56a89ba9cb26e4168cbc37dbd60a28d82430a9e18ad1d25fe0883c314ca2", size = 912322, upload-time = "2026-05-09T23:13:14.281Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/30/e1/c93444052cf41581f3c884ab3fb5823daf0992f11cd4388d4275ca610558/regex-2026.5.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6d189041f15691cfa2b6c4290448ec221244d225b3f5fe9e7771b34ffcdf6e2", size = 801269, upload-time = "2026-05-09T23:13:16.569Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/50/fe/0cf96b882f540e62e8b9956599798203d599c44cf4c77917ca27400ff69b/regex-2026.5.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e82db382b44d0111b22601c509c89f64434816c9e0eef9d1989cda8cc6ff1c04", size = 777085, upload-time = "2026-05-09T23:13:18.675Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/23/5c/d78d4924e7fc875557b9e9b768423925fdfaac5549d06da7810019a9bd26/regex-2026.5.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2acfb48634f64996b57f90f39afa692ff362162722581921fe92239a59960f3c", size = 785153, upload-time = "2026-05-09T23:13:20.525Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bf/e0/5214774090e7b4524dcea3e3c4aa74141d43043f8beb49c1599db1c8b53a/regex-2026.5.9-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:d29eebfc9525db68cad3c97eedd7f754fa265aa5cd0cf4f863b2421e1b48fc9f", size = 860164, upload-time = "2026-05-09T23:13:22.263Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6e/e1/4a57a83350319b1271f0d7a249b8672513ed928b237a741631270de6caea/regex-2026.5.9-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:debb893095e944091c16e641a6e33c1b0f4cb61ab945ec5afbf53ce7068834d8", size = 765731, upload-time = "2026-05-09T23:13:24.277Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/12/f4/499e74a20c156fc75836ee04a72a38d1a063978f600937f9760467beb1b0/regex-2026.5.9-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:d659eee77986549c9ea45b861c7567e44d6287c3dc9a4565478853f7b9fe2ff6", size = 852062, upload-time = "2026-05-09T23:13:26.125Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5b/92/7eebc0d0a01e78629695f342ba17e0deaff8fb45e79cc0d7b98287da6e3e/regex-2026.5.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2efa205e6d98b24d1f3ab395c11aa15cdf10935bca283d0285e0499c284fba21", size = 789577, upload-time = "2026-05-09T23:13:27.814Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/a4/018e71f7d2ad48c1ebe6d3ae0026f9b7cb4802fd15c7cc02fdf724355102/regex-2026.5.9-cp313-cp313-win32.whl", hash = "sha256:f3844f134e834076677dd369976e9f5068679fcb8e50102fdf6b7ac96a3ec127", size = 266691, upload-time = "2026-05-09T23:13:29.549Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e6/1d/861a93719fb9ee7dbfc3761b3797b7a3e112a5d42c6129459d2d741be9b5/regex-2026.5.9-cp313-cp313-win_amd64.whl", hash = "sha256:3527bb4942d2c14552155406cdedd906567456821848aed1cb4933a391bf5eca", size = 277747, upload-time = "2026-05-09T23:13:31.859Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d9/c6/0a2436ae4da1ba76e51cb98943c6838a9a721faa40ebe2dce07694ae34e3/regex-2026.5.9-cp313-cp313-win_arm64.whl", hash = "sha256:56a33f191f17d8c417f99945ebdc1e691d3af9605d86ec68c7e54a57e3e17af6", size = 270500, upload-time = "2026-05-09T23:13:33.525Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/e9/d21346f7b60ed58789371358ed66b09d00f832e1bd7c06e55d9da5679882/regex-2026.5.9-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:01f28d868834624c934b8d2e0aa1c8341337e37831f4a012f18a5afcba4cbaf3", size = 494172, upload-time = "2026-05-09T23:13:35.935Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c4/43/fd1177a2032037c681baecdb3422ee4e1424aec4e4f470ef47793d325274/regex-2026.5.9-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:48036f6374aaa79eb3b754ec29c61d1c6b1606749d705a13f8854fa2539671f6", size = 293952, upload-time = "2026-05-09T23:13:38.307Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f2/7d/9fbf919768368d3f8a4f6c692cf2aa61e482b2b81ec6a298ace4cbf02480/regex-2026.5.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:b96350aa424e79d4fd6b567b344dcbe2b2d6bfc48dfe7717587e1fa6d43da6ff", size = 292314, upload-time = "2026-05-09T23:13:40.353Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e2/6c/e41bfeecb589716843e7c4df09ba46ff2a42961457afece19059d85caeef/regex-2026.5.9-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8f3af7a4903c5c04a11a196a5aa75cdd7dd3f8508132f9fb3259d9f5908e3b88", size = 811681, upload-time = "2026-05-09T23:13:42.543Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/87/83/a5c1c525fba0aa656e88ad0face0b1829788ef4c2fb6b26df58aa1151b84/regex-2026.5.9-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7e87577720152d2caae19fe2baaf1f8d5ca12091e9e229f03915c37d1e4b9178", size = 871135, upload-time = "2026-05-09T23:13:44.326Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/d4/80882e799e440dd878b0979cbebf8fa4d54624a332c83037c7a701649e3f/regex-2026.5.9-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c8b9b9d294cfea3cd19c718ade7cc93492b2c4991abd9a68d0b3477ae6d8e100", size = 917265, upload-time = "2026-05-09T23:13:47.295Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ae/ff/8db60211e2286e396aad7dc7725356c502bff0901ea05bd6cdc2e1a042b9/regex-2026.5.9-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:728d8bfd28a8845c8b6bc5dc7ce010453d206396786c0765c2740cb65f37791e", size = 816311, upload-time = "2026-05-09T23:13:49.885Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4c/47/742ef579c61730f8d268e5cf1f9ce0e37e2ea041ad0f5644724f2378e463/regex-2026.5.9-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7e30b874d341fac767d7df5a0870540541c2c054b80cfaac116e8d367a8a7ff2", size = 785498, upload-time = "2026-05-09T23:13:52.25Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7f/ab/cb0999802dcb0fb95b1ab005e8d4163d8afdd67efc2cb6b6630ac13f8cb1/regex-2026.5.9-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:fd190e88a895a8901325fad284a3f74ea52b1da8525b76cc811fa9b1edf0ce2b", size = 801348, upload-time = "2026-05-09T23:13:54.127Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7d/62/8ca59a24c55bc34d166eefaf3717bd77772f329fdbf984d86581e0a3571c/regex-2026.5.9-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:8e76e8161ad00694cfce6767d5dea860c6391ac5b83e5c3a39661e696f11fc7e", size = 866493, upload-time = "2026-05-09T23:13:56.067Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8d/3d/30f2ae62cef3278bb5bb821f467277a55fb73f01032cf85997e15e8289a8/regex-2026.5.9-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:ddda5340e6c01a293027dd46232fa79eaff1b48058ce7a98f572b6445b088041", size = 772811, upload-time = "2026-05-09T23:13:57.867Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d8/ae/7d2089bcd78ad0c0161bc684339df50032acb438a7bd3305e7ddb1193cec/regex-2026.5.9-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:205109e96b3cf5adf8f4cd62bedde9487feb282b9497a3535451e5a24cd706a0", size = 856584, upload-time = "2026-05-09T23:13:59.679Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a9/29/92ff47f75990131ea4f24ba17819e5a9d141e10819807e09addd73409af6/regex-2026.5.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dfbe4579b9f08036aa7d101d1835437a20783574ac66327e6b29b4018a138081", size = 803453, upload-time = "2026-05-09T23:14:01.978Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/04/99/eff29f1037dcab36702c9ee5d6858cf1ce2336ea8ea2987f64245b99ea5e/regex-2026.5.9-cp313-cp313t-win32.whl", hash = "sha256:ed2c9e8068b614c574d8d30e543d617cf5379b0535d46f97ef00e904745a08b5", size = 269951, upload-time = "2026-05-09T23:14:03.661Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0e/9d/8870b8981d27b22cda77bb26a5ac7ebfa9c7d9e0dea195a834a82380e748/regex-2026.5.9-cp313-cp313t-win_amd64.whl", hash = "sha256:b46b0f094dc1d3b90356c85a0bd2c9bafc4a6a190b9d6f8ddd5a033b6e088ed4", size = 281240, upload-time = "2026-05-09T23:14:05.56Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/b1/3379415e8f135c13ac551353397cc4fe97b4978f3cac73c5fcbcded548b8/regex-2026.5.9-cp313-cp313t-win_arm64.whl", hash = "sha256:872acc074bd29ffc9913ecdfedf6ea77502312ca44a4aa0d3779089c6069d8de", size = 272383, upload-time = "2026-05-09T23:14:07.843Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/13/3e/9c3cd292d8808b3645a2ce517e200179b6d0e903f176300bd8b542e14de5/regex-2026.5.9-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:1bd7587a2948b4085195d5a3374eaf4a425dc3e55784c038175355ecf3bbbf8a", size = 490376, upload-time = "2026-05-09T23:14:09.64Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/70/d43ee8a2ca0a8b68d167f21658b85520ac0574617c7f320367c5047f7556/regex-2026.5.9-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:dea2e88e1cce4522496cce630e11e67b98b7076620bc4336c3f674bc21a375f4", size = 291964, upload-time = "2026-05-09T23:14:11.424Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/21/91/9d50b433828d8e74196904e168a43abf1e6e88b2a15d47ed742456720c37/regex-2026.5.9-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2099f7e7ff7b6aa3192312650a56e91cc091e49d50b04e4f6f8b6e28b3b27f1c", size = 289682, upload-time = "2026-05-09T23:14:13.123Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3e/d2/b835e3cafbb9d977736912436259ff551d60919f7d7b3d37d46659c63564/regex-2026.5.9-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ecd353045824e4477562a2ac718c25799cdaaa41f7aa925a806a8a3e6848a5b9", size = 796996, upload-time = "2026-05-09T23:14:14.923Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2c/a6/9f992d00019166b9de01c546dd4549bc679f2a68df11b877740b0760b7c2/regex-2026.5.9-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:65c8c8c37377794bd5b2f3ebe51919042bf17aec802e23c833d89782ed0c78af", size = 866089, upload-time = "2026-05-09T23:14:17.757Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e0/08/4d32af657e049b19cb62b02e46e38fe1518797bfb2203ee93a510b21b0dc/regex-2026.5.9-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5b73ab8afcf66c622db143d1c6fda4e58e4d537ee4f125229ad47b1ab80f34c0", size = 911530, upload-time = "2026-05-09T23:14:20.353Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d9/27/2af43dd1dc201d1fecefda64a45f4ad0995855b92724f795a777b402ee69/regex-2026.5.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0de5cf193997384ed2ca6f1cd4f78055b255d93d82d5a8cd6ba0d11c10b167e4", size = 800643, upload-time = "2026-05-09T23:14:22.265Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a4/dd/23a249047013b5321d4a60c4d2437462086f601b061776a525e5fba2a59f/regex-2026.5.9-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d641a8c9a61618047796d572a39a79b26167b0411d2c3031937b2fe2d081e2cf", size = 777223, upload-time = "2026-05-09T23:14:24.179Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/94/6a/e85ed9538cd19586d0465076a4578a12e093ce776d15f3f8ce92733a8dd6/regex-2026.5.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:24b2355ef5cc9aa5b8f07d17704face1c166fdcc2290fa7bd6e6c925655a8346", size = 785760, upload-time = "2026-05-09T23:14:26.065Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2a/c4/f25473209438638e947c55f9156fd8f236f74169229028cc99116380868e/regex-2026.5.9-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:a24852d3c29ad9e47593593d8a247c44ccc3d0548ef12c822d6ed0810affe676", size = 860891, upload-time = "2026-05-09T23:14:28.17Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f9/f7/f4f86e3c74419c37370e91f150ae0c2ef7d34b2e0e4cdd5da046a02e4022/regex-2026.5.9-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:916714069da19329ef7de197dcbc77bb3104145c7c2c864dbfbe318f46b88b14", size = 765891, upload-time = "2026-05-09T23:14:30.06Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/26/70/704d8e13765939146b1cd0ef4e2feb71d7929727d2290f026eed10095955/regex-2026.5.9-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:fa411799ca8da32a8d38d020a88faa5b6f91657d284761352940ecf9f7c3bbdd", size = 851380, upload-time = "2026-05-09T23:14:32.123Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/26/29/1a13582a8460038edc38e49f64ceb0dd7c60f5caba77571f4bf6601965d9/regex-2026.5.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:1e6da47d679b7010ef27556b6e0f99771b744936db1792a10ceac6547ae1503e", size = 789350, upload-time = "2026-05-09T23:14:34.799Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/56/3dcafe34fc72e271d62ad9a291801e88a1457bb251c132f15fcc2e5aad1a/regex-2026.5.9-cp314-cp314-win32.whl", hash = "sha256:98bd73080e8756255137e1bd3f3f00295bbc5aa383c0e0f973920e9134d7c4ad", size = 272130, upload-time = "2026-05-09T23:14:36.729Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d0/9c/02eebf0be95efe416c664db7fb8b6b05b7a0b06a7544f2884f2558b0526f/regex-2026.5.9-cp314-cp314-win_amd64.whl", hash = "sha256:ff8d372ac2acdc048d1c19916f27ee61bc5722728458ba6ca5052f2c72d51763", size = 280999, upload-time = "2026-05-09T23:14:39.126Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/70/5a/1dd1abee76cb7a846a0bcf42fdc87e5720c3c33c24f3e37814310a513d9f/regex-2026.5.9-cp314-cp314-win_arm64.whl", hash = "sha256:e1d93bf647916292e8edcec150c07ddf3dc50179ccaf770c04a7f9e452155372", size = 273500, upload-time = "2026-05-09T23:14:41.059Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/86/c1/c5f619b0057a7965cb78ec559c1d7a45ce8c99a35bea95483d64959a93d9/regex-2026.5.9-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:83d0ee4a57d1c87cb549e195ec300b8f0ec3a82eba66d835e4e2ed8634fe4499", size = 494269, upload-time = "2026-05-09T23:14:42.869Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/2c/5d01f1aee33de4bbe60c8452945bfc8477ca7c5ae4450f6bfe711036cb36/regex-2026.5.9-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:d3d7eb5c9a7f6df82ed3cfac9beb93882a5cbcb5b8b157b56cb2b3b276574ac1", size = 293954, upload-time = "2026-05-09T23:14:44.822Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7a/fe/e8988b2ae2108c6ef71bd4aa8d87fbe257976dd0810e826cd75f701c68b6/regex-2026.5.9-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:075160bf16658e16d35233300b8453aac25de4cbea808d22348b6979668e924d", size = 292405, upload-time = "2026-05-09T23:14:47.211Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/79/34/d2b0937faa7859263f7f0a3c6b103a1296306be6952dc173d0154e9a2f49/regex-2026.5.9-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:45375819235558a4ff1c4971dc32881f022613abdb180128f5cb4768c1765a1c", size = 811855, upload-time = "2026-05-09T23:14:49.21Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/80/fe/daf53a47457a8486db66c66c01ceb9c2303eecee3f87197f1e77eb1a736d/regex-2026.5.9-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ead4b163ac30a29574510cd4b3e2e985ac5290c05fc7095557d6a5f403fc31b5", size = 871189, upload-time = "2026-05-09T23:14:51.555Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1c/75/058fc4470cbfbf57d800aff1a0022b929a3f9fa553ee10a0cdf2070eb31f/regex-2026.5.9-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8c6e4218fbdfbcd4f6c19efca40930d24a621bf4b48cb76bc6640543bd28ef20", size = 917485, upload-time = "2026-05-09T23:14:53.633Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/88/e7/179cfda3a28bc843b5c6cfe7f79f23489c791ed95f151083803660878432/regex-2026.5.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6351571c8a42b505eb555c0dc47d740d0fb66977dc142919eea6f4325b7c56a0", size = 816369, upload-time = "2026-05-09T23:14:56.198Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/41/90/6f0cc422071688266d344fca8462d787cba0a2c144acb25721f9a61ec265/regex-2026.5.9-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:002205cafd2a9e78c6290c7d1df277bf3277b3b7a30e0b4bb0dac2e2e3f7cb2d", size = 785869, upload-time = "2026-05-09T23:14:58.602Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/02/67/a31f1760f09c27b251ef39e9beb541f462cf977381d067faa764c2c0e393/regex-2026.5.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8abd33fef90b2a9efac5557d6033ca82d1195ed3a15fea5af15ba7b463c6a63b", size = 801427, upload-time = "2026-05-09T23:15:00.642Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e3/c4/1a80654597b6bc1e1ea0494824c31200e8a956abe290afae9b19a166a148/regex-2026.5.9-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:31037c82eccb44b7ea2e9e221d7c01429430e989a1f4b91ea5a855f6017b509a", size = 866482, upload-time = "2026-05-09T23:15:03.384Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d1/11/960724e06482c08466ff5611e242e86f80062949cdf6b4b9cc317b9dd93d/regex-2026.5.9-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:5604dfd046dc37eca90250fc3be938b076c8059fa772ac0ed6f499b0f0fb0415", size = 773022, upload-time = "2026-05-09T23:15:05.625Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/50/a8/a9979c3e7918280e93159ebcab5ef1a65116dd4f3bd6091be0eae4a126e8/regex-2026.5.9-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:0e1b1b4e496afbb24f4a62aba855ee4f88f25578927697b340702e48c9ee6bc2", size = 856642, upload-time = "2026-05-09T23:15:07.966Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fe/d4/a9b732f2f0072c0ab12227483abb24fffcb9f73f8a2b203df0a6d0434735/regex-2026.5.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:be3372b9df6ddecff6486d37e19095a7b4973137caf5512407a89f4455361f41", size = 803552, upload-time = "2026-05-09T23:15:10.215Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d5/fe/1b3113817447a1d4155e4ac76d2e072f42c0bcba2f43fa8a0e756ea2cd91/regex-2026.5.9-cp314-cp314t-win32.whl", hash = "sha256:3ddd90103f9e5c471c49c7852ecc1fe27c7e45eb99e977aefe7caa4e779f4f58", size = 275746, upload-time = "2026-05-09T23:15:12.609Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/92/73/93d42045302636c91f2e5ef588b65b84b01428f28ec77de256b1dfdfbe5c/regex-2026.5.9-cp314-cp314t-win_amd64.whl", hash = "sha256:ca518ed29c46eecba6010b15f1b9a479314d2de409536e71b6a13aa04e3b8a77", size = 285685, upload-time = "2026-05-09T23:15:15.086Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/80/35b4c33c804a165a7f55289afda3ea9e3eb6d15800341a2d66455c0f1f30/regex-2026.5.9-cp314-cp314t-win_arm64.whl", hash = "sha256:5e41809d2683fcde7d5a8c87a6567ba1fb1ce0de9f31bff578de00a4b2d76daa", size = 275713, upload-time = "2026-05-09T23:15:16.98Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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.7.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, +] + +[[package]] +name = "sentencepiece" +version = "0.2.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d8/15/46afbab00733d81788b64be430ca1b93011bb9388527958e26cc31832de5/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6356d0986b8b8dc351b943150fcd81a1c6e6e4d439772e8584c64230e58ca987", size = 1942560, upload-time = "2025-08-12T06:59:25.82Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fa/79/7c01b8ef98a0567e9d84a4e7a910f8e7074fcbf398a5cd76f93f4b9316f9/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:8f8ba89a3acb3dc1ae90f65ec1894b0b9596fdb98ab003ff38e058f898b39bc7", size = 1325385, upload-time = "2025-08-12T06:59:27.722Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bb/88/2b41e07bd24f33dcf2f18ec3b74247aa4af3526bad8907b8727ea3caba03/sentencepiece-0.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:02593eca45440ef39247cee8c47322a34bdcc1d8ae83ad28ba5a899a2cf8d79a", size = 1253319, upload-time = "2025-08-12T06:59:29.306Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a0/54/38a1af0c6210a3c6f95aa46d23d6640636d020fba7135cd0d9a84ada05a7/sentencepiece-0.2.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a0d15781a171d188b661ae4bde1d998c303f6bd8621498c50c671bd45a4798e", size = 1316162, upload-time = "2025-08-12T06:59:30.914Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ef/66/fb191403ade791ad2c3c1e72fe8413e63781b08cfa3aa4c9dfc536d6e795/sentencepiece-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4f5a3e0d9f445ed9d66c0fec47d4b23d12cfc858b407a03c194c1b26c2ac2a63", size = 1387785, upload-time = "2025-08-12T06:59:32.491Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a9/2d/3bd9b08e70067b2124518b308db6a84a4f8901cc8a4317e2e4288cdd9b4d/sentencepiece-0.2.1-cp311-cp311-win32.whl", hash = "sha256:6d297a1748d429ba8534eebe5535448d78b8acc32d00a29b49acf28102eeb094", size = 999555, upload-time = "2025-08-12T06:59:34.475Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/32/b8/f709977f5fda195ae1ea24f24e7c581163b6f142b1005bc3d0bbfe4d7082/sentencepiece-0.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:82d9ead6591015f009cb1be1cb1c015d5e6f04046dbb8c9588b931e869a29728", size = 1054617, upload-time = "2025-08-12T06:59:36.461Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7a/40/a1fc23be23067da0f703709797b464e8a30a1c78cc8a687120cd58d4d509/sentencepiece-0.2.1-cp311-cp311-win_arm64.whl", hash = "sha256:39f8651bd10974eafb9834ce30d9bcf5b73e1fc798a7f7d2528f9820ca86e119", size = 1033877, upload-time = "2025-08-12T06:59:38.391Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/89/fa/d3d5ebcba3cb9e6d3775a096251860c41a6bc53a1b9461151df83fe93255/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99f955df238021bf11f0fc37cdb54fd5e5b5f7fd30ecc3d93fb48b6815437167", size = 1316273, upload-time = "2025-08-12T06:59:44.476Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/04/88/14f2f4a2b922d8b39be45bf63d79e6cd3a9b2f248b2fcb98a69b12af12f5/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdfecef430d985f1c2bcbfff3defd1d95dae876fbd0173376012d2d7d24044b", size = 1387881, upload-time = "2025-08-12T06:59:46.09Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fd/b8/903e5ccb77b4ef140605d5d71b4f9e0ad95d456d6184688073ed11712809/sentencepiece-0.2.1-cp312-cp312-win32.whl", hash = "sha256:a483fd29a34c3e34c39ac5556b0a90942bec253d260235729e50976f5dba1068", size = 999540, upload-time = "2025-08-12T06:59:48.023Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2d/81/92df5673c067148c2545b1bfe49adfd775bcc3a169a047f5a0e6575ddaca/sentencepiece-0.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:4cdc7c36234fda305e85c32949c5211faaf8dd886096c7cea289ddc12a2d02de", size = 1054671, upload-time = "2025-08-12T06:59:49.895Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fe/02/c5e3bc518655d714622bec87d83db9cdba1cd0619a4a04e2109751c4f47f/sentencepiece-0.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:daeb5e9e9fcad012324807856113708614d534f596d5008638eb9b40112cd9e4", size = 1033923, upload-time = "2025-08-12T06:59:51.952Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ba/4a/85fbe1706d4d04a7e826b53f327c4b80f849cf1c7b7c5e31a20a97d8f28b/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dcd8161eee7b41aae57ded06272905dbd680a0a04b91edd0f64790c796b2f706", size = 1943150, upload-time = "2025-08-12T06:59:53.588Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c2/83/4cfb393e287509fc2155480b9d184706ef8d9fa8cbf5505d02a5792bf220/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c6c8f42949f419ff8c7e9960dbadcfbc982d7b5efc2f6748210d3dd53a7de062", size = 1325651, upload-time = "2025-08-12T06:59:55.073Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8d/de/5a007fb53b1ab0aafc69d11a5a3dd72a289d5a3e78dcf2c3a3d9b14ffe93/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:097f3394e99456e9e4efba1737c3749d7e23563dd1588ce71a3d007f25475fff", size = 1253641, upload-time = "2025-08-12T06:59:56.562Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2c/d2/f552be5928105588f4f4d66ee37dd4c61460d8097e62d0e2e0eec41bc61d/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d7b670879c370d350557edabadbad1f6561a9e6968126e6debca4029e5547820", size = 1316271, upload-time = "2025-08-12T06:59:58.109Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/96/df/0cfe748ace5485be740fed9476dee7877f109da32ed0d280312c94ec259f/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c7f0fd2f2693309e6628aeeb2e2faf6edd221134dfccac3308ca0de01f8dab47", size = 1387882, upload-time = "2025-08-12T07:00:00.701Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ac/dd/f7774d42a881ced8e1739f393ab1e82ece39fc9abd4779e28050c2e975b5/sentencepiece-0.2.1-cp313-cp313-win32.whl", hash = "sha256:92b3816aa2339355fda2c8c4e021a5de92180b00aaccaf5e2808972e77a4b22f", size = 999541, upload-time = "2025-08-12T07:00:02.709Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dd/e9/932b9eae6fd7019548321eee1ab8d5e3b3d1294df9d9a0c9ac517c7b636d/sentencepiece-0.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:10ed3dab2044c47f7a2e7b4969b0c430420cdd45735d78c8f853191fa0e3148b", size = 1054669, upload-time = "2025-08-12T07:00:04.915Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c9/3a/76488a00ea7d6931689cda28726a1447d66bf1a4837943489314593d5596/sentencepiece-0.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:ac650534e2251083c5f75dde4ff28896ce7c8904133dc8fef42780f4d5588fcd", size = 1033922, upload-time = "2025-08-12T07:00:06.496Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4a/b6/08fe2ce819e02ccb0296f4843e3f195764ce9829cbda61b7513f29b95718/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8dd4b477a7b069648d19363aad0cab9bad2f4e83b2d179be668efa672500dc94", size = 1946052, upload-time = "2025-08-12T07:00:08.136Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ab/d9/1ea0e740591ff4c6fc2b6eb1d7510d02f3fb885093f19b2f3abd1363b402/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0c0f672da370cc490e4c59d89e12289778310a0e71d176c541e4834759e1ae07", size = 1327408, upload-time = "2025-08-12T07:00:09.572Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/99/7e/1fb26e8a21613f6200e1ab88824d5d203714162cf2883248b517deb500b7/sentencepiece-0.2.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad8493bea8432dae8d6830365352350f3b4144415a1d09c4c8cb8d30cf3b6c3c", size = 1254857, upload-time = "2025-08-12T07:00:11.021Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bc/85/c72fd1f3c7a6010544d6ae07f8ddb38b5e2a7e33bd4318f87266c0bbafbf/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b81a24733726e3678d2db63619acc5a8dccd074f7aa7a54ecd5ca33ca6d2d596", size = 1315722, upload-time = "2025-08-12T07:00:12.989Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4a/e8/661e5bd82a8aa641fd6c1020bd0e890ef73230a2b7215ddf9c8cd8e941c2/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0a81799d0a68d618e89063fb423c3001a034c893069135ffe51fee439ae474d6", size = 1387452, upload-time = "2025-08-12T07:00:15.088Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/99/5e/ae66c361023a470afcbc1fbb8da722c72ea678a2fcd9a18f1a12598c7501/sentencepiece-0.2.1-cp313-cp313t-win32.whl", hash = "sha256:89a3ea015517c42c0341d0d962f3e6aaf2cf10d71b1932d475c44ba48d00aa2b", size = 1002501, upload-time = "2025-08-12T07:00:16.966Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c1/03/d332828c4ff764e16c1b56c2c8f9a33488bbe796b53fb6b9c4205ddbf167/sentencepiece-0.2.1-cp313-cp313t-win_amd64.whl", hash = "sha256:33f068c9382dc2e7c228eedfd8163b52baa86bb92f50d0488bf2b7da7032e484", size = 1057555, upload-time = "2025-08-12T07:00:18.573Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/88/14/5aee0bf0864df9bd82bd59e7711362908e4935e3f9cdc1f57246b5d5c9b9/sentencepiece-0.2.1-cp313-cp313t-win_arm64.whl", hash = "sha256:b3616ad246f360e52c85781e47682d31abfb6554c779e42b65333d4b5f44ecc0", size = 1036042, upload-time = "2025-08-12T07:00:20.209Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/9c/89eb8b2052f720a612478baf11c8227dcf1dc28cd4ea4c0c19506b5af2a2/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:5d0350b686c320068702116276cfb26c066dc7e65cfef173980b11bb4d606719", size = 1943147, upload-time = "2025-08-12T07:00:21.809Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/82/0b/a1432bc87f97c2ace36386ca23e8bd3b91fb40581b5e6148d24b24186419/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c7f54a31cde6fa5cb030370566f68152a742f433f8d2be458463d06c208aef33", size = 1325624, upload-time = "2025-08-12T07:00:23.289Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ea/99/bbe054ebb5a5039457c590e0a4156ed073fb0fe9ce4f7523404dd5b37463/sentencepiece-0.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c83b85ab2d6576607f31df77ff86f28182be4a8de6d175d2c33ca609925f5da1", size = 1253670, upload-time = "2025-08-12T07:00:24.69Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/ad/d5c7075f701bd97971d7c2ac2904f227566f51ef0838dfbdfdccb58cd212/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1855f57db07b51fb51ed6c9c452f570624d2b169b36f0f79ef71a6e6c618cd8b", size = 1316247, upload-time = "2025-08-12T07:00:26.435Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/03/35fbe5f3d9a7435eebd0b473e09584bd3cc354ce118b960445b060d33781/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01e6912125cb45d3792f530a4d38f8e21bf884d6b4d4ade1b2de5cf7a8d2a52b", size = 1387894, upload-time = "2025-08-12T07:00:28.339Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dc/aa/956ef729aafb6c8f9c443104c9636489093bb5c61d6b90fc27aa1a865574/sentencepiece-0.2.1-cp314-cp314-win32.whl", hash = "sha256:c415c9de1447e0a74ae3fdb2e52f967cb544113a3a5ce3a194df185cbc1f962f", size = 1096698, upload-time = "2025-08-12T07:00:29.764Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b8/cb/fe400d8836952cc535c81a0ce47dc6875160e5fedb71d2d9ff0e9894c2a6/sentencepiece-0.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:881b2e44b14fc19feade3cbed314be37de639fc415375cefaa5bc81a4be137fd", size = 1155115, upload-time = "2025-08-12T07:00:32.865Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/32/89/047921cf70f36c7b6b6390876b2399b3633ab73b8d0cb857e5a964238941/sentencepiece-0.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:2005242a16d2dc3ac5fe18aa7667549134d37854823df4c4db244752453b78a8", size = 1133890, upload-time = "2025-08-12T07:00:34.763Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a1/11/5b414b9fae6255b5fb1e22e2ed3dc3a72d3a694e5703910e640ac78346bb/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a19adcec27c524cb7069a1c741060add95f942d1cbf7ad0d104dffa0a7d28a2b", size = 1946081, upload-time = "2025-08-12T07:00:36.97Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/77/eb/7a5682bb25824db8545f8e5662e7f3e32d72a508fdce086029d89695106b/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:e37e4b4c4a11662b5db521def4e44d4d30ae69a1743241412a93ae40fdcab4bb", size = 1327406, upload-time = "2025-08-12T07:00:38.669Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/03/b0/811dae8fb9f2784e138785d481469788f2e0d0c109c5737372454415f55f/sentencepiece-0.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:477c81505db072b3ab627e7eab972ea1025331bd3a92bacbf798df2b75ea86ec", size = 1254846, upload-time = "2025-08-12T07:00:40.611Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ef/23/195b2e7ec85ebb6a547969f60b723c7aca5a75800ece6cc3f41da872d14e/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:010f025a544ef770bb395091d57cb94deb9652d8972e0d09f71d85d5a0816c8c", size = 1315721, upload-time = "2025-08-12T07:00:42.914Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7e/aa/553dbe4178b5f23eb28e59393dddd64186178b56b81d9b8d5c3ff1c28395/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:733e59ff1794d26db706cd41fc2d7ca5f6c64a820709cb801dc0ea31780d64ab", size = 1387458, upload-time = "2025-08-12T07:00:44.56Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/66/7c/08ff0012507297a4dd74a5420fdc0eb9e3e80f4e88cab1538d7f28db303d/sentencepiece-0.2.1-cp314-cp314t-win32.whl", hash = "sha256:d3233770f78e637dc8b1fda2cd7c3b99ec77e7505041934188a4e7fe751de3b0", size = 1099765, upload-time = "2025-08-12T07:00:46.058Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/91/d5/2a69e1ce15881beb9ddfc7e3f998322f5cedcd5e4d244cb74dade9441663/sentencepiece-0.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:5e4366c97b68218fd30ea72d70c525e6e78a6c0a88650f57ac4c43c63b234a9d", size = 1157807, upload-time = "2025-08-12T07:00:47.673Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f3/16/54f611fcfc2d1c46cbe3ec4169780b2cfa7cf63708ef2b71611136db7513/sentencepiece-0.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:105e36e75cbac1292642045458e8da677b2342dcd33df503e640f0b457cb6751", size = 1136264, upload-time = "2025-08-12T07:00:49.485Z" }, +] + +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/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://pypi.tuna.tsinghua.edu.cn/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 = "sparkle-engine" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "fastapi" }, + { name = "mlx" }, + { name = "mlx-lm" }, + { name = "numpy" }, + { name = "textual" }, + { name = "uvicorn" }, +] + +[package.dev-dependencies] +dev = [ + { name = "httpx" }, + { name = "nanobind" }, + { name = "pytest" }, + { name = "pytest-asyncio" }, +] + +[package.metadata] +requires-dist = [ + { name = "fastapi", specifier = ">=0.115" }, + { name = "mlx", specifier = ">=0.31" }, + { name = "mlx-lm", specifier = ">=0.31" }, + { name = "numpy", specifier = ">=2.0" }, + { name = "textual", specifier = ">=0.80" }, + { name = "uvicorn", specifier = ">=0.30" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "httpx", specifier = ">=0.27" }, + { name = "nanobind", specifier = ">=2.12.0" }, + { name = "pytest", specifier = ">=8" }, + { name = "pytest-asyncio", specifier = ">=0.23" }, +] + +[[package]] +name = "starlette" +version = "1.3.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "anyio" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/eb/e3/7c1dc7381d9f8ab7d854328ebfa884e62cb3f3d8549ddfd37c7814f42afa/starlette-1.3.1.tar.gz", hash = "sha256:05d0213193f2fbaae60e2ecb593b4add4262ad4e46536b54abe36f11a71724e0", size = 2703240, upload-time = "2026-06-12T09:23:11.602Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ec/bb/2799cc2ede3ed41131f8975621e7213dfc7ef4acbbaadfa440f32500c370/starlette-1.3.1-py3-none-any.whl", hash = "sha256:c7372aae11c3c3f26a42df7bd626cec2f47d03483d261d369516a615a53714c6", size = 73632, upload-time = "2026-06-12T09:23:10.017Z" }, +] + +[[package]] +name = "textual" +version = "8.2.8" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "markdown-it-py", extra = ["linkify"] }, + { name = "mdit-py-plugins" }, + { name = "platformdirs" }, + { name = "pygments" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/00/21/39a76b01bd5eea82a04baaca7580e105d8c59450df03998345bb2cfb307b/textual-8.2.8.tar.gz", hash = "sha256:3f106a9fbc73e39dd266c9712432087de78a6d644084c7c241d6a25c3169115b", size = 1860502, upload-time = "2026-06-30T06:51:24.495Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/be/35261223d9416a0751cdff1c7b4a6f881387218a12d439fe22fefebc8c04/textual-8.2.8-py3-none-any.whl", hash = "sha256:267375fd402dc8d981457212efa71f0e3365fd17bba144ba9bb3ed7563cb374a", size = 731418, upload-time = "2026-06-30T06:51:26.364Z" }, +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "huggingface-hub" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/92/97/5dbfabf04c7e348e655e907ed27913e03db0923abb5dfdd120d7b25630e1/tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c", size = 3100275, upload-time = "2026-01-05T10:41:02.158Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2e/47/174dca0502ef88b28f1c9e06b73ce33500eedfac7a7692108aec220464e7/tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001", size = 2981472, upload-time = "2026-01-05T10:41:00.276Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/78/59/09d0d9ba94dcd5f4f1368d4858d24546b4bdc0231c2354aa31d6199f0399/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd", size = 3168835, upload-time = "2026-01-05T10:40:38.847Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/47/50/b3ebb4243e7160bda8d34b731e54dd8ab8b133e50775872e7a434e524c28/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5", size = 3521673, upload-time = "2026-01-05T10:40:56.614Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e0/fa/89f4cb9e08df770b57adb96f8cbb7e22695a4cb6c2bd5f0c4f0ebcf33b66/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e", size = 3724818, upload-time = "2026-01-05T10:40:44.507Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/64/04/ca2363f0bfbe3b3d36e95bf67e56a4c88c8e3362b658e616d1ac185d47f2/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b", size = 3379195, upload-time = "2026-01-05T10:40:51.139Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6c/fb/66e2da4704d6aadebf8cb39f1d6d1957df667ab24cff2326b77cda0dcb85/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a", size = 9560069, upload-time = "2026-01-05T10:45:10.673Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/16/04/fed398b05caa87ce9b1a1bb5166645e38196081b225059a6edaff6440fac/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a", size = 9899263, upload-time = "2026-01-05T10:45:12.559Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fd/18/a545c4ea42af3df6effd7d13d250ba77a0a86fb20393143bbb9a92e434d4/tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92", size = 2502363, upload-time = "2026-01-05T10:45:20.593Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/65/71/0670843133a43d43070abeb1949abfdef12a86d490bea9cd9e18e37c5ff7/tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48", size = 2747786, upload-time = "2026-01-05T10:45:18.411Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, +] + +[[package]] +name = "tqdm" +version = "4.68.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/06/b3/36c8ecf72e8925200671613332db156d84b99b3aee742a41c1938ebb0808/tqdm-4.68.1.tar.gz", hash = "sha256:fc163d96b287bd031e1aa24421ce4411b25559bd0a1be4fe649bdaa4d2c02bf5", size = 171236, upload-time = "2026-06-05T17:23:15.267Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/47/aa/218a0eb34de1f753c83e4d0d1c8e7c4cef27f20dcb8342e024f63a80dc86/tqdm-4.68.1-py3-none-any.whl", hash = "sha256:fea4a90e4023f764914569f7802a297277c5ab1a66be5144143e142e1a4031d8", size = 78354, upload-time = "2026-06-05T17:23:13.654Z" }, +] + +[[package]] +name = "transformers" +version = "5.10.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8d/38/d5f978bd5091019e89aef29b9a831f5cd70f2598963a3ead8b9570cab592/transformers-5.10.2.tar.gz", hash = "sha256:f9a44b9c8ca9ab1156b467f574d832ea066284299c2fd0ed84641ccb592751fc", size = 8799687, upload-time = "2026-06-04T18:43:49.119Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/6f/e1564b0cc182afa05e219a8e09a8e770ffaab879b6b824b56c819bd221da/transformers-5.10.2-py3-none-any.whl", hash = "sha256:8a669db546f82c7c3618cb46ceb0f0afd89292bc70f319c058f8332ec63e268d", size = 11003830, upload-time = "2026-06-04T18:43:45.303Z" }, +] + +[[package]] +name = "typer" +version = "0.25.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "click" }, + { name = "rich" }, + { name = "shellingham" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e4/51/9aed62104cea109b820bbd6c14245af756112017d309da813ef107d42e7e/typer-0.25.1.tar.gz", hash = "sha256:9616eb8853a09ffeabab1698952f33c6f29ffdbceb4eaeecf571880e8d7664cc", size = 122276, upload-time = "2026-04-30T19:32:16.964Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3f/f9/2b3ff4e56e5fa7debfaf9eb135d0da96f3e9a1d5b27222223c7296336e5f/typer-0.25.1-py3-none-any.whl", hash = "sha256:75caa44ed46a03fb2dab8808753ffacdbfea88495e74c85a28c5eefcf5f39c89", size = 58409, upload-time = "2026-04-30T19:32:18.271Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "uc-micro-py" +version = "2.0.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/78/67/9a363818028526e2d4579334460df777115bdec1bb77c08f9db88f6389f2/uc_micro_py-2.0.0.tar.gz", hash = "sha256:c53691e495c8db60e16ffc4861a35469b0ba0821fe409a8a7a0a71864d33a811", size = 6611, upload-time = "2026-03-01T06:31:27.526Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/61/73/d21edf5b204d1467e06500080a50f79d49ef2b997c79123a536d4a17d97c/uc_micro_py-2.0.0-py3-none-any.whl", hash = "sha256:3603a3859af53e5a39bc7677713c78ea6589ff188d70f4fee165db88e22b242c", size = 6383, upload-time = "2026-03-01T06:31:26.257Z" }, +] + +[[package]] +name = "uvicorn" +version = "0.51.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "click" }, + { name = "h11" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a2/65/b7c6c443ccc58678c91e1e973bbe2a878591538655d6e1d47f24ba1c51f3/uvicorn-0.51.0.tar.gz", hash = "sha256:f6f4b69b657c312f516dd2d268ab9ae6f254b11e4bac504f37b2ab58b24dd0b0", size = 94412, upload-time = "2026-07-08T10:59:05.962Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/45/ec/dbb7e5a6b91f86bfb9eb7d2988a2730907b6a729875b949c7f022e8b88fa/uvicorn-0.51.0-py3-none-any.whl", hash = "sha256:5d38af6cd620f2ae3849fb44fd4879e0890aa1febe8d47eb355fb45d93fe6a5b", size = 73219, upload-time = "2026-07-08T10:59:04.44Z" }, +]