jungege112358/sparkle
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Browse Source
This commit is contained in:
fiser_jun 2026-08-04 14:34:00 +08:00
commit 4745f264b2
90 changed files with 13695 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
.gitignore vendored Normal file
View File

@ -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

187
README.md Normal file
View File

@ -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: 回答里夹「<E5A4B9>」乱码？**
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 性能基线 |

61
benchmarks/reports/8bit-g64-baseline-2026-07-25.md Normal file
View File

@ -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)

3
mlx_streaming/.gitignore vendored Normal file
View File

@ -0,0 +1,3 @@
.venv/
__pycache__/
*.pyc

0
mlx_streaming/__init__.py Normal file
View File

329
mlx_streaming/cli.py Normal file
View File

@ -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())

298
mlx_streaming/config.py Normal file
View File

@ -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) # 切点：层号 <cutoff 用 lo，否则 hi
def cross_layer_ahead_lo() -> 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

0
mlx_streaming/core/__init__.py Normal file
View File

1
mlx_streaming/core/cache/__init__.py vendored Normal file
View File

@ -0,0 +1 @@
"""专家缓存模块：常驻池(resident_pool) / LRU+文件后端存储(expert_store) / blob 流式源(blob_loader)。"""

174
mlx_streaming/core/cache/autopin.py vendored Normal file
View File

@ -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<b：以 c 为新基线
self._cpp_baseline = cur
def flush(self) -> 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

238
mlx_streaming/core/cache/blob_loader.py vendored Normal file
View File

@ -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()

604
mlx_streaming/core/cache/expert_store.py vendored Normal file
View File

@ -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("<Q", data[:8])[0]
header = json.loads(data[8:8 + n])
base = 8 + n
out: "Dict[str, mx.array]" = {}
for name, info in header.items():
if name == "__metadata__":
continue
np_dt, view_dt = _ST_DTYPE[info["dtype"]]
s, e = info["data_offsets"]
arr = mx.array(np.frombuffer(data[base + s:base + e], dtype=np_dt).reshape(info["shape"]))
if view_dt is not None:
arr = arr.view(view_dt)
out[name] = arr
return out
class _PerLayerLru:
"""每层一个 OrderedDict 的 LRU 集合，容量为「每层槽数」。"""
def __init__(self, capacity: int, clear_on_evict: bool = False):
self.capacity = capacity
self.clear_on_evict = clear_on_evict
self._caches: "Dict[int, OrderedDict[int, Dict[str, mx.array]]]" = {}
self.hits = 0
self.misses = 0
self.prefetch_hits = 0
self.prefetch_loads = 0
def _layer_cache(self, layer: int) -> "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

116
mlx_streaming/core/cache/kv_quant_patch.py vendored Normal file
View File

@ -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

223
mlx_streaming/core/cache/quant_kv.py vendored Normal file
View File

@ -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

635
mlx_streaming/core/cache/resident_pool.py vendored Normal file
View File

@ -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

237
mlx_streaming/core/cache/virtual_pool.py vendored Normal file
View File

@ -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)

1
mlx_streaming/core/linear_attn/__init__.py Normal file
View File

@ -0,0 +1 @@
"""线性注意力（Qwen3-Next gated-delta）相关的自定义 kernel。"""

226
mlx_streaming/core/linear_attn/gated_delta_multistate.py Normal file
View File

@ -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<float>(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<InT>(out);
}}
}} else {{
y[dv_idx] = static_cast<InT>(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<StT>(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<StT>(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)

75
mlx_streaming/core/mem.py Normal file
View File

@ -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

1
mlx_streaming/core/moe/__init__.py Normal file
View File

@ -0,0 +1 @@
"""MoE 模块：门控选专家(gate) / 专家计算(compute) / 自定义算子(custom_kernel) / 热路径块(block)。"""

431
mlx_streaming/core/moe/block.py Normal file
View File

@ -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

208
mlx_streaming/core/moe/compute.py Normal file
View File

@ -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)

266
mlx_streaming/core/moe/custom_kernel.py Normal file
View File

@ -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

32
mlx_streaming/core/moe/gate.py Normal file
View File

@ -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

504
mlx_streaming/core/moe/native_moe.py Normal file
View File

@ -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

1
mlx_streaming/core/prefetch/__init__.py Normal file
View File

@ -0,0 +1 @@
"""专家预取模块：后台物化(bg_prefetch) / native staging(native_staging) / 跨层预测预取(cross_layer) / 模型 patch(patch)。"""

133
mlx_streaming/core/prefetch/bg_prefetch.py Normal file
View File

@ -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)

181
mlx_streaming/core/prefetch/cross_layer.py Normal file
View File

@ -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

229
mlx_streaming/core/prefetch/native_staging.py Normal file
View File

@ -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

56
mlx_streaming/core/prefetch/patch.py Normal file
View File

@ -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

131
mlx_streaming/core/profiling.py Normal file
View File

@ -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

50
mlx_streaming/core/route_trace.py Normal file
View File

@ -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)

252
mlx_streaming/model_builder.py Normal file
View File

@ -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)

0
mlx_streaming/models/__init__.py Normal file
View File

0
mlx_streaming/mtp/__init__.py Normal file
View File

28
mlx_streaming/mtp/bench_verdict.py Normal file
View File

@ -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"

188
mlx_streaming/mtp/drafter.py Normal file
View File

@ -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)

506
mlx_streaming/mtp/generate.py Normal file
View File

@ -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

276
mlx_streaming/mtp/kv_cache.py Normal file
View File

@ -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

95
mlx_streaming/mtp/qwen3_next_mtp.py Normal file
View File

@ -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

BIN
mlx_streaming/native_moe_ext.cpython-314-darwin.so Executable file
View File

Binary file not shown.

0
mlx_streaming/prep/__init__.py Normal file
View File

32
mlx_streaming/prep/blob_layout.py Normal file
View File

@ -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

35
mlx_streaming/prep/download_retry.py Normal file
View File

@ -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()

131
mlx_streaming/prep/extract_mtp.py Normal file
View File

@ -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()

87
mlx_streaming/prep/pack_blob_from_experts.py Normal file
View File

@ -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()

83
mlx_streaming/prep/pack_compute_buffers.py Normal file
View File

@ -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()

70
mlx_streaming/prep/pack_expert_bundles.py Normal file
View File

@ -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()

135
mlx_streaming/prep/pack_expert_ranges.py Normal file
View File

@ -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("<Q", f.read(8))[0]
header = json.loads(f.read(header_len))
base = 8 + header_len
out = []
for key, meta in header.items():
if key == "__metadata__":
continue
start, end = meta["data_offsets"]
f.seek(base + start)
payload = f.read(end - start)
out.append({
"key": key,
"dtype": meta["dtype"],
"shape": meta["shape"],
"payload": payload,
})
return out
def pack_layer(src_dir: str, out_dir: str, layer: int, num_experts: int) -> 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()

92
mlx_streaming/prep/repack_expert_blobs.py Normal file
View File

@ -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()

69
mlx_streaming/prep/split_experts.py Normal file
View File

@ -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))

1
mlx_streaming/runtime/__init__.py Normal file
View File

@ -0,0 +1 @@
"""生产推理入口：基线/流式/投机解码等 run_* 启动脚本。"""

42
mlx_streaming/runtime/run_baseline.py Normal file
View File

@ -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()

69
mlx_streaming/runtime/run_mmap.py Normal file
View File

@ -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()

409
mlx_streaming/runtime/run_mtp_spec.py Normal file
View File

@ -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()

116
mlx_streaming/runtime/run_spec.py Normal file
View File

@ -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()

180
mlx_streaming/runtime/run_streaming.py Normal file
View File

@ -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()

9
mlx_streaming/server/__init__.py Normal file
View File

@ -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"]

95
mlx_streaming/server/__main__.py Normal file
View File

@ -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()

116
mlx_streaming/server/admin.py Normal file
View File

@ -0,0 +1,116 @@
"""极简 admin 页:单文件 HTML + 内联 JS,不引框架/CDN。中文文案,2.5s 轮询 stats。"""
ADMIN_HTML = """<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Sparkle 引擎控制台</title>
<style>
body { font-family: -apple-system, "PingFang SC", sans-serif; max-width: 640px;
margin: 32px auto; padding: 0 16px; color: #222; }
h1 { font-size: 20px; }
section { border: 1px solid #ddd; border-radius: 8px; padding: 16px;
margin-bottom: 16px; }
.row { margin: 8px 0; }
.label { display: inline-block; width: 180px; color: #555; }
input[type=number] { width: 90px; }
input[type=range] { width: 240px; vertical-align: middle; }
button { padding: 6px 20px; }
#msg { margin-left: 12px; color: #060; }
.hint { color: #888; font-size: 13px; }
</style>
</head>
<body>
<h1>Sparkle 引擎控制台</h1>
<section>
<div class="row"><span class="label">引擎状态</span><span id="loaded">-</span></div>
<div class="row"><span class="label">模型</span><span id="model">-</span></div>
<div class="row"><span class="label">最近一轮吞吐</span><span id="tps">-</span> tok/s</div>
<div class="row"><span class="label">进程内存 RSS</span><span id="rss">-</span></div>
</section>
<section>
<div class="row">
<span class="label">常驻专家池 expert_slots</span>
<input type="range" id="slots" min="32" max="160" step="1" value="32">
<b id="slotsVal">32</b>
</div>
<div class="row hint">预估池内存:<span id="poolMem"></span>(slots × ~160MB(g64));修改后引擎后台重建,期间对话不可用</div>
<div class="row">
<span class="label">投机宽度 k</span>
<input type="number" id="k" min="1" max="8" step="1" value="3">
</div>
<div class="row">
<span class="label">最大生成 max_tokens</span>
<input type="number" id="maxTokens" min="1" max="32768" step="1" value="4096">
</div>
<div class="row">
<button id="apply">应用</button><span id="msg"></span>
</div>
</section>
<script>
const $ = id => document.getElementById(id);
function fmtBytes(b) {
if (b >= 1024 * 1024 * 1024) return (b / 1024 / 1024 / 1024).toFixed(2) + " GB";
return (b / 1024 / 1024).toFixed(0) + " MB";
}
function updateSlots() {
const v = +$("slots").value;
$("slotsVal").textContent = v;
$("poolMem").textContent = " " + fmtBytes(v * 160 * 1024 * 1024) + " ";
}
$("slots").addEventListener("input", updateSlots);
async function refresh() {
try {
const s = await (await fetch("/api/stats")).json();
$("loaded").textContent = s.loaded ? "已加载" : "未加载 / 重建中";
$("model").textContent = s.model || "-";
$("tps").textContent = (s.tok_per_s || 0).toFixed(1);
$("rss").textContent = fmtBytes(s.rss_bytes || 0);
if (document.activeElement !== $("slots")) {
$("slots").value = s.expert_slots;
updateSlots();
}
} catch (e) { /* 服务暂不可达时下轮再试 */ }
}
$("apply").addEventListener("click", async () => {
const body = {
expert_slots: +$("slots").value,
k: +$("k").value,
max_tokens: +$("maxTokens").value
};
try {
const r = await (await fetch("/api/engine/config", {
method: "POST",
headers: {"Content-Type": "application/json"},
body: JSON.stringify(body)
})).json();
$("msg").textContent = r.reloading ? "已应用,引擎正在后台重建…" : "已应用";
} catch (e) {
$("msg").textContent = "应用失败";
}
setTimeout(() => { $("msg").textContent = ""; }, 4000);
});
(async function init() {
try {
const c = await (await fetch("/api/engine/config")).json();
$("slots").value = c.expert_slots;
$("k").value = c.k;
$("maxTokens").value = c.max_tokens;
} catch (e) { /* 用默认值 */ }
updateSlots();
refresh();
setInterval(refresh, 2500);
})();
</script>
</body>
</html>
"""

104
mlx_streaming/server/app.py Normal file
View File

@ -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

338
mlx_streaming/server/openai.py Normal file
View File

@ -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
会把它们渲染成模型训练时学过的格式(<tool_call> 与 <tool_response> 包裹)。
此前把工具返回改写成普通 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":
# 原生透传:模板渲染为 <tool_response>,模型按训练格式 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"):
# 原生透传:模板渲染为 <tool_call>。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 模板把工具定义注入提示词),模型据此输出
# <tool_call> 文本,客户端兜底解析——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

101
mlx_streaming/server/state.py Normal file
View File

@ -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

9
mlx_streaming/tui/__init__.py Normal file
View File

@ -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

334
mlx_streaming/tui/app.py Normal file
View File

@ -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)

436
mlx_streaming/tui/backend.py Normal file
View File

@ -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 被真字符替换后长度不变 → 增量为空、真字符永久丢失,客户端就留下一堆「<E5A086>」。
# detokenizer 的 .text 只包含已完整的字符(未完成字节留在 _unflushed),恒为下一步的真前缀。
detok = self._detok
if detok is None:
detok = self._detok = tok.detokenizer
detok.reset()
produced_all: list[int] = [] # 已上报的正文 token(不含 EOS 及其后)
stopped = {"v": False}
hit_eos = {"v": False}
def on_tokens(new_ids):
for t in new_ids:
if int(t) in eos:
hit_eos["v"] = True
break
detok.add_token(int(t))
produced_all.append(int(t))
if hit_eos["v"]:
detok.finalize() # 冲出尾部残留字节(正常收尾,不会有半字)
if on_text(detok.text, len(produced_all)): # 用户按 Esc 请求中断
stopped["v"] = True
return True
# 命中 EOS:完整回答已生成,提前停止,避免引擎空跑到 max_tokens 让界面
# 长时间卡在「思考中」。EOS 属正常完成,不算中断。
return hit_eos["v"]
t0 = time.perf_counter()
if sampling is not None:
# 采样模式(server 默认,质量优先):非投机逐 token 采样,无 MTP 加速。
produced, forwarded = self._generate_sampled(
ids, on_tokens, self.args.max_tokens, sampling,
main_cache, cached_len)
stats = {"resident_tokens": len(ids) + forwarded}
else:
produced, stats = mtp_generate(
self._model,
self._drafter,
tok,
mx.array([ids]),
self.args.max_tokens,
K=self.args.k,
ids_mode=True,
profile=False,
on_tokens=on_tokens,
main_cache=main_cache,
cached_len=cached_len,
)
dt = time.perf_counter() - t0
# 持久化本轮 cache 供下轮复用。正常情况下 main_cache 恰好持有 `ids + produced[:-1]`
# (produced[-1] 为 pending 未入 cache)。但末步多 token 跨 max_tokens 会 over-commit:
# cache 领先于 produced,无法用已知 token 精确表述——此时禁用复用,下轮全量重建,绝不错算。
resident = stats.get("resident_tokens")
expected = len(ids) + len(produced) - 1
if resident == expected:
self._main_cache = main_cache
self._cached_ids = list(ids) + list(produced[:-1])
else:
self._main_cache = None
self._cached_ids = []
out_ids = _truncate_eos(produced, eos)
text = tok.decode(out_ids)
tps = len(out_ids) / dt if dt > 0 else 0.0
return GenResult(text, len(out_ids), tps, stopped=stopped["v"])

3
mlx_streaming/tui/banner.py Normal file
View File

@ -0,0 +1,3 @@
"""sparkle 顶栏 logo。保持单行,便于与模型信息拼在同一行。"""
LOGO = "▚ sparkle"

41
mlx_streaming/tui/styles.tcss Normal file
View File

@ -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;
}

49
models/README.md Normal file
View File

@ -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/` 会在首次跑通后自动生成，可不预置。

32
native/ext/CMakeLists.txt Normal file
View File

@ -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()

14
native/ext/Makefile Normal file
View File

@ -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

153
native/ext/bindings.cpp Normal file
View File

@ -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<nb::gil_scoped_release>()); // 等待时释放 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<nb::gil_scoped_release>()); // 阻塞等时释放 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());
}

34
native/ext/common.h Normal file
View File

@ -0,0 +1,34 @@
// 共用前导：系统/Metal/MLX/nanobind 头 + 命名空间别名。所有 TU 都包含。
#pragma once
#include <nanobind/nanobind.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/variant.h>
#include <nanobind/stl/vector.h>
#include <Metal/Metal.hpp>
#include <algorithm>
#include <atomic>
#include <cctype>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <fcntl.h>
#include <map>
#include <mutex>
#include <random>
#include <utility>
#include <stdexcept>
#include <string>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <vector>
#include "mlx/backend/metal/device.h"
#include "mlx/mlx.h"
#include "mlx/primitives.h"
namespace nb = nanobind;
namespace mx = mlx::core;

664
native/ext/compute/fused_moe.cpp Normal file
View File

@ -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<size_t>(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<size_t>(out_dim) * words * sizeof(uint32_t),
static_cast<size_t>(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<uint16_t>(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<uint32_t*>(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<uint16_t*>(sbuf->contents());
auto* b = reinterpret_cast<uint16_t*>(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<int>& ids,
MTL::Buffer* wbuf,
MTL::Buffer* sbuf,
MTL::Buffer* bbuf) {
auto* wd = static_cast<char*>(wbuf->contents());
auto* sd = static_cast<char*>(sbuf->contents());
auto* bd = static_cast<char*>(bbuf->contents());
auto* ws = static_cast<const char*>(src.w);
auto* ss = static_cast<const char*>(src.s);
auto* bs = static_cast<const char*>(src.b);
for (size_t local = 0; local < ids.size(); ++local) {
size_t expert = static_cast<size_t>(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 <metal_stdlib>
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<float>(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<int> 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<mx::array>&, std::vector<mx::array>&) override {
throw std::runtime_error("FusedMoePrimitive only supports GPU evaluation");
}
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
auto& x = inputs[0];
auto& scores = inputs[1];
auto& out = outputs[0];
int active = static_cast<int>(expert_ids_.size());
int tokens = static_cast<int>(x.size() / hidden_);
int k = active / std::max(1, tokens);
if (active <= 0 || tokens <= 0 || active != static_cast<int>(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<size_t>(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<size_t>(active) * gu.weight_per;
size_t gu_s_bytes = static_cast<size_t>(active) * gu.scale_per;
size_t down_w_bytes = static_cast<size_t>(active) * down.weight_per;
size_t down_s_bytes = static_cast<size_t>(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<size_t>(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<int> expert_ids_;
MappedProj gate_;
MappedProj up_;
MappedProj down_;
std::vector<NS::SharedPtr<MTL::Buffer>> 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<mx::array>&, std::vector<mx::array>&) override {
throw std::runtime_error("StagedFusedMoePrimitive only supports GPU evaluation");
}
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
auto& x = inputs[0];
auto& scores = inputs[1];
auto& out = outputs[0];
int tokens = static_cast<int>(x.size() / hidden_);
int active = static_cast<int>(scores.size());
int k = active / std::max(1, tokens);
if (active <= 0 || tokens <= 0 || active != static_cast<int>(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<size_t>(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<size_t>(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<mx::array>&, std::vector<mx::array>&) override {
throw std::runtime_error("SlotFusedMoePrimitive only supports GPU evaluation");
}
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
auto& x = inputs[0];
auto& local_slots = inputs[1];
auto& scores = inputs[2];
auto& out = outputs[0];
int tokens = static_cast<int>(x.size() / hidden_);
int active = static_cast<int>(scores.size());
int k = active / std::max(1, tokens);
if (active <= 0 || tokens <= 0 || active != static_cast<int>(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<size_t>(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<size_t>(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<int> expert_vec(ids.size());
const uint32_t* idp = ids.data<uint32_t>();
for (size_t i = 0; i < ids.size(); ++i) expert_vec[i] = static_cast<int>(idp[i]);
auto out_shape = x.shape();
out_shape.back() = hidden;
return mx::array(
out_shape,
mx::float32,
std::make_shared<FusedMoePrimitive>(
mx::to_stream(s),
compute_dir,
layer,
hidden,
inter,
group,
bits,
num_experts,
synthetic,
std::move(expert_vec)),
std::vector<mx::array>{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<StagedFusedMoePrimitive>(mx::to_stream(s), hidden, inter, group, bits),
std::vector<mx::array>{
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<SlotFusedMoePrimitive>(mx::to_stream(s), hidden, inter, group, bits),
std::vector<mx::array>{
x, local_slots, scores,
gate_w, gate_s, gate_b,
up_w, up_s, up_b,
down_w, down_s, down_b});
}

23
native/ext/compute/fused_moe.h Normal file
View File

@ -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);

183
native/ext/io/bg_reader.cpp Normal file
View File

@ -0,0 +1,183 @@
// [7] 自由后台读线程（de-risk）：双队列 + 低优并发上限的后台 pread 线程池。
#include "bg_reader.h"
#include "blob_io.h"
#include <condition_variable>
#include <functional>
#include <queue>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <fcntl.h>
#include <unistd.h>
namespace {
struct ReadOp { uint8_t* dst; size_t nbytes; off_t file_off; };
struct BgJob {
std::vector<ReadOp> ops;
std::string path;
long ticket;
std::vector<mx::array> keep; // 持 buffer 引用，保证 dst 指针在读期间存活
std::function<void()> 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<std::mutex> 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<std::mutex> 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<std::mutex> 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<void()> fn) {
ensure_started(4);
BgJob job;
job.task = std::move(fn);
{ std::lock_guard<std::mutex> lk(m_); low_q_.push(std::move(job)); } // 通用任务走低优
cv_.notify_all();
}
bool ready(long ticket) {
std::lock_guard<std::mutex> lk(dm_);
return done_.count(ticket) > 0;
}
void wait(long ticket) {
std::unique_lock<std::mutex> lk(dm_);
dcv_.wait(lk, [&] { return done_.count(ticket) > 0; });
}
void stop() {
{ std::lock_guard<std::mutex> lk(m_); running_ = false; }
cv_.notify_all();
for (auto& t : threads_) if (t.joinable()) t.join();
threads_.clear();
{ std::lock_guard<std::mutex> lk(dm_); done_.clear(); }
{ std::lock_guard<std::mutex> 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<std::string, int> fds; // 每线程各自缓存 fd
while (true) {
std::unique_lock<std::mutex> 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<ssize_t>(op.nbytes))
fprintf(stderr, "[bg pread SHORT] got=%zd want=%zu off=%lld\n",
got, op.nbytes, static_cast<long long>(op.file_off));
}
{ std::lock_guard<std::mutex> lk2(dm_); done_.insert(job.ticket); }
dcv_.notify_all();
}
if (is_low) { // 释放低优额度 → 唤醒别的 worker 再取
std::lock_guard<std::mutex> 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<BgJob> high_q_, low_q_;
int active_low_ = 0; // 当前正在执行的低优读数
int low_cap_ = 0; // 低优并发上限（<=0 不限流，保持旧行为）
std::unordered_set<long> done_;
std::vector<std::thread> threads_;
bool running_ = false;
};
BgReader g_bg;
} // namespace
void bg_submit_task(std::function<void()> 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<int>& experts,
const std::vector<int>& rows, const std::string& path,
int stride, long ticket, int prio) {
mx::array d = dst;
d.eval();
uint8_t* base = d.data<uint8_t>();
size_t st = static_cast<size_t>(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<size_t>(rows[i]) * st, st,
static_cast<off_t>(static_cast<size_t>(experts[i]) * st)});
g_bg.submit(std::move(job));
return ticket;
}
// 专家各段直写进多个池段张量的 slot 行（消费侧零 MLX 算子）。
long bg_pread_into_pool(
const std::vector<mx::array>& dst,
const std::vector<long>& seg_off,
const std::vector<long>& 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<uint8_t>();
job.keep.push_back(d);
job.ops.push_back(ReadOp{
base + static_cast<size_t>(slot) * static_cast<size_t>(seg_nb[i]),
static_cast<size_t>(seg_nb[i]),
static_cast<off_t>(static_cast<size_t>(expert) * static_cast<size_t>(stride)
+ static_cast<size_t>(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(); }

23
native/ext/io/bg_reader.h Normal file
View File

@ -0,0 +1,23 @@
// 自由后台读线程（de-risk）：脱离 GPU 完成回调、零 GIL，pread 进调用方 MLX buffer。
// 线程全程只碰 dst 原始指针 / 路径 / 整数，绝不接触 Python 对象 → 无需 GIL。
#pragma once
#include "../common.h"
#include <functional>
void bg_reader_start(int workers, int low_cap = 0);
long bg_reader_submit(const mx::array& dst, const std::vector<int>& experts,
const std::vector<int>& 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<mx::array>& dst,
const std::vector<long>& seg_off,
const std::vector<long>& 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<void()> fn);

16
native/ext/io/blob_io.h Normal file
View File

@ -0,0 +1,16 @@
// 共用底层 IO 助手：打开 blob 文件。多个模块（侧区/后台读线程等）的 pread 都用它。
#pragma once
#include <fcntl.h>
#include <unistd.h>
#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;
}

58
native/ext/io/blob_load.cpp Normal file
View File

@ -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 <fcntl.h>
#include <unistd.h>
class BlobLoadPrimitive : public mx::Primitive {
public:
BlobLoadPrimitive(mx::Stream stream, std::string path, size_t stride, std::vector<int> 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<mx::array>&, std::vector<mx::array>& outputs) override {
load(outputs[0]);
}
void eval_gpu(const std::vector<mx::array>&, std::vector<mx::array>& outputs) override {
load(outputs[0]);
}
private:
void load(mx::array& out) {
out.set_data(mx::allocator::malloc(out.nbytes()));
uint8_t* dst = out.data<uint8_t>();
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<size_t>(experts_[i]) * stride_;
ssize_t n = ::pread(fd, dst + i * stride_, stride_, static_cast<off_t>(off));
if (n != static_cast<ssize_t>(stride_)) {
::close(fd);
throw std::runtime_error("blob pread short read");
}
}
::close(fd);
}
std::string path_;
size_t stride_;
std::vector<int> 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<int> ev(ids.size());
const uint32_t* p = ids.data<uint32_t>();
for (size_t i = 0; i < ids.size(); ++i) ev[i] = static_cast<int>(p[i]);
int n = static_cast<int>(ev.size());
return mx::array(
mx::Shape{n, stride},
mx::uint8,
std::make_shared<BlobLoadPrimitive>(mx::to_stream(s), path, static_cast<size_t>(stride), std::move(ev)),
std::vector<mx::array>{});
}

7
native/ext/io/blob_load.h Normal file
View File

@ -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);

331
native/ext/pool/demand.cpp Normal file
View File

@ -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 <chrono>
#include <cstdio>
#include <cstdlib>
#include <unordered_map>
#include <unordered_set>
struct RealLayer {
std::vector<int> order; // 插入序(LRU tie-break)，与 e2r 同步维护
std::unordered_map<int, int> e2r; // expert -> slot [0,cap)
std::vector<int> free_rows; // 空闲槽(front pop，仿 free.pop(0))
std::unordered_map<int, uint32_t> freq; // LFU 频次(驱逐不删，与 Python 一致)
std::unordered_set<int> pinned; // AUTOPIN 钉死集：驱逐永不选(见 real_pin)
int cap = 0;
long access = 0; // 累计访问(decay 用)
bool inited = false;
};
static std::mutex g_real_mutex;
static std::map<int, RealLayer> 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<long> 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<double, std::micro>(
std::chrono::steady_clock::now().time_since_epoch()).count();
}
std::vector<double> 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<std::mutex> lk(g_real_mutex);
real_ensure_locked(g_real[layer], cap);
}
std::vector<int> real_region_contents(int layer) {
std::lock_guard<std::mutex> lk(g_real_mutex);
std::vector<int> 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<std::mutex> lk(g_real_mutex);
auto it = g_real.find(layer);
return it == g_real.end() ? 0 : static_cast<int>(it->second.e2r.size());
}
int real_pinned_count(int layer) {
std::lock_guard<std::mutex> lk(g_real_mutex);
auto it = g_real.find(layer);
return it == g_real.end() ? 0 : static_cast<int>(it->second.pinned.size());
}
void real_reset() {
std::lock_guard<std::mutex> 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<int>& 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<int> real_pin(int layer, const std::vector<int>& experts) {
std::lock_guard<std::mutex> lk(g_real_mutex);
RealLayer& c = g_real[layer];
std::vector<int> 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<long> real_freq_dump() {
std::lock_guard<std::mutex> lk(g_real_mutex);
std::vector<long> out;
for (auto& kv : g_real)
for (auto& p : kv.second.freq) {
out.push_back(static_cast<long>(kv.first));
out.push_back(static_cast<long>(p.first));
out.push_back(static_cast<long>(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<int>& 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<int>& uniq_access,
bool lfu, int decay_interval) {
if (!lfu) return;
for (int e : uniq_access) c.freq[e] += 1;
c.access += static_cast<long>(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<int32_t> demand_core_locked(
RealLayer& c, const uint32_t* ip, size_t n, const std::unordered_map<int, int>& side,
bool lfu, int decay_interval, std::vector<std::pair<int, int>>& placements, long stats[4]) {
std::vector<int32_t> local(n, -1);
std::vector<int> uniq_miss, access_order;
std::unordered_set<int> miss_seen, access_seen;
int hitpos = 0;
// pass1：算命中(侧区覆盖真实区)、收集 miss(首见序)、收集唯一访问(freq)。
for (size_t i = 0; i < n; ++i) {
int e = static_cast<int>(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<int, int> 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<int>(ip[i]));
local[i] = (it != new_slot.end()) ? it->second : 0;
++misspos;
}
}
stats[0] = hitpos; stats[1] = misspos; stats[2] = static_cast<long>(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<mx::array>& pool_list,
const std::vector<int>& 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<uint32_t>();
double t1 = g_dt_on ? dt_now_us() : 0;
std::vector<uint8_t*> ptrs;
ptrs.reserve(pool_list.size());
for (auto& a0 : pool_list) { mx::array a = a0; a.eval(); ptrs.push_back(a.data<uint8_t>()); }
double t2 = g_dt_on ? dt_now_us() : 0;
std::unordered_map<int, int> side = sideregion_snapshot(layer, side_gen); // 侧区快照(该代)
double ta = g_dt_on ? dt_now_us() : 0;
long stats[4];
std::vector<int32_t> local;
std::vector<std::pair<int, int>> placements; // (expert, slot)：锁外并行落盘
double tb;
{
std::lock_guard<std::mutex> 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<long> 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<long> 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<long>(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<std::mutex> 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<long> 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<long> demand_last_stats() {
std::lock_guard<std::mutex> 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<int> real_debug_place(int layer, const std::vector<int>& experts_flat, int cap,
bool lfu, int decay_interval) {
std::vector<uint32_t> u(experts_flat.begin(), experts_flat.end());
long stats[4];
std::vector<std::pair<int, int>> placements;
std::lock_guard<std::mutex> lk(g_real_mutex);
RealLayer& c = g_real[layer];
real_ensure_locked(c, cap);
std::unordered_map<int, int> empty_side;
auto local = demand_core_locked(c, u.data(), u.size(), empty_side, lfu, decay_interval,
placements, stats);
return std::vector<int>(local.begin(), local.end());
}

29
native/ext/pool/demand.h Normal file
View File

@ -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<int> real_region_contents(int layer); // [expert, slot, ...]
int real_region_count(int layer);
void real_reset();
// AUTOPIN：注册 pinned 热专家(驱逐免选)，返回平行槽位(-1=分不到)；前置 real_init。
std::vector<int> real_pin(int layer, const std::vector<int>& experts);
// 该层 pinned(永不可驱逐)专家数。调用方据此算「可安放的唯一专家上限」= cap − pinned，
// 用于在超容量前把该前向分流到 host/fetch 路径（超了会落 0 号槽 → 逐位错算）。
int real_pinned_count(int layer);
// AUTOPIN：导出各层 LFU 累计频次，扁平 [layer, expert, count, ...]。
std::vector<long> real_freq_dump();
// demand 全接管：inds 惰性(内部 eval 一次)；side_gen 侧区代；pool_list 为 _segs 顺序池数组。
mx::array demand_dual(
const mx::array& inds, const std::vector<mx::array>& pool_list,
const std::vector<int>& 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<long> demand_last_stats();
std::vector<double> demand_timings(); // [inds_eval, pool_eval, state, build] us
void demand_timing_enable(bool on);
// 测试壳：纯状态推进(不 pread/不侧区)，返回 local 槽位；供 LFU 驱逐等价对拍。
std::vector<int> real_debug_place(int layer, const std::vector<int>& experts_flat, int cap,
bool lfu, int decay_interval);

89
native/ext/pool/owned_pool.cpp Normal file
View File

@ -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<mx::allocator::Buffer> 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<int>& 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<size_t>(d);
size_t nbytes = n * static_cast<size_t>(mx::size_of(dt));
auto buf = mx::allocator::malloc(nbytes);
std::memset(buf.raw_ptr(), 0, nbytes);
{
std::lock_guard<std::mutex> 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<mx::array>& pool_list,
const std::vector<mx::array>& srcs_flat,
const std::vector<int>& slots) {
int K = static_cast<int>(pool_list.size());
int m = static_cast<int>(slots.size());
if (m == 0) return;
if (static_cast<int>(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<uint8_t>();
for (int j = 0; j < m; ++j) {
mx::array s = srcs_flat[static_cast<size_t>(i) * m + j];
s.eval();
size_t nb = s.nbytes();
std::memcpy(base + static_cast<size_t>(slots[j]) * nb, s.data<uint8_t>(), nb);
}
}
}
// 同上，但源是每 key 预堆叠的 (m,*shape) 整块 stacked_list[i]，按行 memcpy 到 slots[j]。
void pool_write_stacked(const std::vector<mx::array>& pool_list,
const std::vector<mx::array>& stacked_list,
const std::vector<int>& slots) {
int K = static_cast<int>(pool_list.size());
int m = static_cast<int>(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<uint8_t>();
const uint8_t* src = st.data<uint8_t>();
size_t rownb = st.nbytes() / static_cast<size_t>(m); // stacked 为 (m,*shape) → 每行字节
for (int j = 0; j < m; ++j)
std::memcpy(base + static_cast<size_t>(slots[j]) * rownb,
src + static_cast<size_t>(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<uintptr_t>(b.data<uint8_t>());
}

14
native/ext/pool/owned_pool.h Normal file
View File

@ -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<int>& shape, const std::string& dtype);
// demand 真实区落池：把已加载专家段直接 memcpy 进 owned 池行（无 MLX scatter，保 buffer 稳定）。
void pool_write_rows(const std::vector<mx::array>& pool_list,
const std::vector<mx::array>& srcs_flat, const std::vector<int>& slots);
void pool_write_stacked(const std::vector<mx::array>& pool_list,
const std::vector<mx::array>& stacked_list, const std::vector<int>& slots);
// 诊断：mx.array 底层 buffer 原始指针（uintptr），用于对拍侧区写入 buffer 与 consume 读到 buffer 是否同一块。
uintptr_t array_data_ptr(const mx::array& a);

389
native/ext/pool/side_region.cpp Normal file
View File

@ -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 <condition_variable>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <thread>
#include <unordered_set>
struct SideLayer {
std::map<int, int> e2r; // expert -> 物理侧区行 [base_row, base_row+spec)
std::vector<int> free_rows;
std::map<int, uint32_t> 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<std::pair<int, int>, SideLayer> g_side; // 键 (layer, gen)：双缓冲两代独立
// 侧区 fill 在途计数：eval_gpu 提交预取时 +1，后台 read_publish 写完字节后 -1。
// 消费方在前向开头 sideregion_drain() 排空上一前向的 fill，保证被消费的侧区行字节
// 已完全写好（含 GPU 完成回调滞后的情形）→ 消灭「GPU 消费 kernel 读到半写侧区行」竞态。
static std::atomic<long> 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<std::mutex> lk(g_side_drain_mutex);
g_side_inflight.fetch_sub(1);
}
g_side_drain_cv.notify_all();
}
void sideregion_drain() {
std::unique_lock<std::mutex> 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<uint64_t> g_side_ev{0};
static inline unsigned side_tid() {
return static_cast<unsigned>(
std::hash<std::thread::id>{}(std::this_thread::get_id()) & 0xffff);
}
class PrefetchPoolSideRegionPrimitive : public mx::Primitive {
public:
PrefetchPoolSideRegionPrimitive(mx::Stream s, std::vector<int> seg_nbytes, int layer, int gen,
std::string path, size_t stride, std::vector<int> 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<mx::array>& in, std::vector<mx::array>& out) override {
out[0].set_data(mx::allocator::malloc(out[0].nbytes()));
run(in);
}
void eval_gpu(const std::vector<mx::array>& in, std::vector<mx::array>& out) override {
out[0].set_data(mx::allocator::malloc(out[0].nbytes()));
std::vector<uint8_t*> 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<uint8_t>());
}
mx::array ids = in[0];
const uint32_t* idp = ids.data<uint32_t>();
size_t n = ids.size();
std::vector<int> seg = seg_;
int layer = layer_;
int gen = gen_;
std::string path = path_;
size_t stride = stride_;
std::vector<int> 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<mx::array>& in) {
std::vector<uint8_t*> 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<uint8_t>());
}
mx::array ids = in[0];
ids.eval();
// CPU 路径同步执行（测试/无 GPU 时）：预留 + 读发布一气呵成。
auto to_read = reserve(ids.data<uint32_t>(), 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<std::pair<int, int>> reserve(
const uint32_t* idp, size_t n, int layer, int gen, const std::vector<int>& 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<int> res(resident.begin(), resident.end());
std::vector<int> P;
std::unordered_set<int> Pset, seen;
for (size_t i = 0; i < n; ++i) {
int e = static_cast<int>(idp[i]);
if (res.count(e) || !seen.insert(e).second) continue;
P.push_back(e);
Pset.insert(e);
}
std::vector<std::pair<int, int>> to_read;
std::lock_guard<std::mutex> 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<std::pair<int, int>>& to_read) {
std::map<int, int> 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<int> 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<uint8_t*>& ptrs, const std::vector<int>& seg,
const std::vector<std::pair<int, int>>& to_read,
const std::string& path, size_t stride, int layer, int gen) {
// 段在 blob 记录内的偏移（段顺序 = seg 顺序 = ptrs/池 key 顺序）。
std::vector<size_t> 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<size_t>(seg[i]); }
int fd = open_blob_nocache(path.c_str());
if (fd < 0) { // 失败则把预留行还回 free
std::lock_guard<std::mutex> 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<uint8_t> rec(stride); // 整条 blob 记录临时缓冲
std::vector<std::pair<int, int>> done;
for (auto& pr : to_read) {
int e = pr.first, row = pr.second;
if (::pread(fd, rec.data(), stride, static_cast<off_t>(static_cast<size_t>(e) * stride)) !=
static_cast<ssize_t>(stride)) {
std::lock_guard<std::mutex> 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<size_t>(row) * static_cast<size_t>(seg[i]),
rec.data() + seg_off[i], static_cast<size_t>(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<std::mutex> 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<int> seg_;
int layer_;
int gen_;
std::string path_;
size_t stride_;
std::vector<int> resident_;
int spec_;
int base_;
};
mx::array prefetch_pool_sideregion(
const std::vector<mx::array>& pool_list, const std::vector<int>& seg_nbytes,
const mx::array& expert_ids, int layer, const std::string& path, int stride,
const std::vector<int>& 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<size_t>(b);
if (seg_sum != static_cast<size_t>(stride)) {
throw std::invalid_argument(
"prefetch_pool_sideregion: sum(seg_nbytes)=" + std::to_string(seg_sum) +
" != stride=" + std::to_string(stride));
}
std::vector<mx::array> 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<PrefetchPoolSideRegionPrimitive>(
mx::to_stream(s), seg_nbytes, layer, gen, path, static_cast<size_t>(stride), resident,
spec_slots, base_row),
inputs);
}
std::vector<int> sideregion_contents(int layer, int gen) {
std::lock_guard<std::mutex> lk(g_side_mutex);
std::vector<int> 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<mx::array, mx::array> sideregion_kv(int layer, int gen) {
std::vector<uint32_t> keys;
std::vector<int32_t> vals;
{
std::lock_guard<std::mutex> 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<uint32_t>(p.first));
vals.push_back(static_cast<int32_t>(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<int>(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<std::mutex> lk(g_side_mutex);
g_side.clear();
}
std::unordered_map<int, int> sideregion_snapshot(int layer, int gen) {
std::unordered_map<int, int> side;
std::lock_guard<std::mutex> 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;
}

21
native/ext/pool/side_region.h Normal file
View File

@ -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 <unordered_map>
mx::array prefetch_pool_sideregion(
const std::vector<mx::array>& pool_list, const std::vector<int>& seg_nbytes,
const mx::array& expert_ids, int layer, const std::string& path, int stride,
const std::vector<int>& resident, int spec_slots, int base_row, int gen,
mx::StreamOrDevice s);
std::vector<int> sideregion_contents(int layer, int gen); // [expert, phys_row, ...]
std::pair<mx::array, mx::array> 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<int, int> sideregion_snapshot(int layer, int gen);

217
native/ext/prefetch/prefetch.cpp Normal file
View File

@ -0,0 +1,217 @@
// [2] 轻量预取(无 staging，仅预热 page cache) + [3] staging 版 miss→hit + STAGING_HPROF 探针。
#include "prefetch.h"
#include "../io/bg_reader.h"
#include <chrono>
#include <functional>
#include <map>
#include <tuple>
#include <unordered_set>
#include <fcntl.h>
#include <unistd.h>
// ====== [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<mx::array>& inputs, std::vector<mx::array>& outputs) override {
outputs[0].set_data(mx::allocator::malloc(outputs[0].nbytes()));
mx::array ids = inputs[0];
ids.eval();
record_and_read(ids.data<uint32_t>(), ids.size(), path_, stride_, do_read_);
}
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& 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<uint32_t>(); // 指针此刻有效（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<uint8_t> buf;
if (do_read && buf.size() < stride) buf.resize(stride);
for (size_t i = 0; i < n; ++i) {
int e = static_cast<int>(p[i]);
if (fd >= 0) ::pread(fd, buf.data(), stride, static_cast<off_t>(static_cast<size_t>(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<PrefetchOnCompletePrimitive>(
mx::to_stream(s), path, static_cast<size_t>(stride), do_read),
std::vector<mx::array>{expert_ids});
}
// ---- handler 触发时刻探针(STAGING_HPROF)：记录每次 staging 完成回调被 Metal 触发的时刻 ----
// 用于实测"完成回调到底扎堆在 eval 尾、还是逐层铺开"。仅诊断用，默认关。
static std::mutex g_hprof_mutex;
static std::vector<std::tuple<long, int, double>> g_hprof_log; // (gen, layer, t_fire_seconds)
static bool g_hprof_on = false;
static inline double hprof_steady_now() {
return std::chrono::duration<double>(
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<int, std::pair<long, std::vector<std::pair<int, int>>>> 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<int> 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<mx::array>& inputs, std::vector<mx::array>& 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<uint32_t>(), ids.size(), stg.data<uint8_t>(), layer_, gen_, path_, stride_,
resident_, cap_);
}
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& 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<uint32_t>();
uint8_t* sp = stg.data<uint8_t>();
size_t n = ids.size();
int layer = layer_; long gen = gen_; std::string path = path_; size_t stride = stride_;
std::vector<int> 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<std::mutex> 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<int>& resident, int cap) {
int fd = ::open(path.c_str(), O_RDONLY);
if (fd < 0) return;
std::unordered_set<int> res(resident.begin(), resident.end());
std::unordered_set<int> done; // 去重：同一缺口只 pread 一次
std::vector<std::pair<int, int>> ready;
ready.reserve(static_cast<size_t>(cap));
int row = 0; // 写入的 staging 行（≤ cap）
for (size_t i = 0; i < n && row < cap; ++i) {
int e = static_cast<int>(idp[i]);
if (res.count(e) || !done.insert(e).second) continue; // 已常驻/已取过 → 跳过
if (::pread(fd, stg + static_cast<size_t>(row) * stride, stride,
static_cast<off_t>(static_cast<size_t>(e) * stride))
== static_cast<ssize_t>(stride)) {
ready.emplace_back(e, row);
++row;
}
}
::close(fd);
std::lock_guard<std::mutex> 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<int> 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<int>& resident, int cap,
bool parallel, mx::StreamOrDevice s = {}) {
return mx::array(
mx::Shape{1}, mx::uint8,
std::make_shared<PrefetchStagingPrimitive>(
mx::to_stream(s), layer, gen, path, static_cast<size_t>(stride), resident, cap, parallel),
std::vector<mx::array>{expert_ids, staging});
}
// 取走某层就绪记录：[gen, e0,r0,e1,r1,...]（首元素是 generation）；空表示无就绪。
std::vector<long> prefetch_staging_take(int layer) {
std::lock_guard<std::mutex> lk(g_stg_mutex);
std::vector<long> 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<std::mutex> 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<double> staging_hprof_get() {
std::lock_guard<std::mutex> lk(g_hprof_mutex);
std::vector<double> out;
out.reserve(g_hprof_log.size() * 3);
for (auto& r : g_hprof_log) {
out.push_back(static_cast<double>(std::get<0>(r)));
out.push_back(static_cast<double>(std::get<1>(r)));
out.push_back(std::get<2>(r));
}
return out;
}

23
native/ext/prefetch/prefetch.h Normal file
View File

@ -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<int>& resident, int cap,
bool parallel, mx::StreamOrDevice s);
// 取走某层就绪记录：[gen, e0,r0,e1,r1,...]；空表示无就绪。
std::vector<long> prefetch_staging_take(int layer);
// handler 触发时刻探针(诊断用)：enable 开关并清零；now 取同时钟当前秒；get 取 (gen,layer,t) 日志。
void staging_hprof_enable(bool on);
double staging_hprof_now();
std::vector<double> staging_hprof_get(); // 扁平 [gen,layer,t, ...]

44
pyproject.toml Normal file
View File

@ -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"

1129
uv.lock generated Normal file
View File

File diff suppressed because it is too large Load Diff