From cad29a1e1d327eb1fba55fb81de07ae507b38616 Mon Sep 17 00:00:00 2001 From: Bin Date: Mon, 28 Sep 2026 16:36:25 +0800 Subject: [PATCH 01/13] refactor(repo): adopt the multi-platform layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prepare the repo for the four-platform story (Python / macOS / iOS / Android) ahead of the unified inference framework: - move the Python framework (pyproject, src, tests, scripts, examples) into python/; add python/README.md (pyproject readme target); docs stay at the repo root - reserve macos/, ios/, android/ platform directories; their READMEs describe the released engines (source drops 2026-09-30) - CI: run both jobs with working-directory python/ - hygiene suite: scan from the repo root so it covers all platform dirs; exclude .pdf; MLX-boundary check follows python/src - fix(edge0 convert-adapters): resolve scripts/ via parents[2] — the command pointed at a non-existent src/scripts since before the move - .gitignore: python/scripts/alignment_results.json Verified: pytest -m 'not slow' → 79 passed, 1 skipped (fresh venv, editable install from python/). Co-Authored-By: Claude Code --- .github/workflows/ci.yml | 6 +++++ .gitignore | 2 +- android/README.md | 23 ++++++++++++++++++ docs/adding-a-model.md | 2 +- docs/prerouter.md | 2 +- ios/README.md | 23 ++++++++++++++++++ macos/README.md | 24 +++++++++++++++++++ python/README.md | 24 +++++++++++++++++++ {examples => python/examples}/bench.py | 0 {examples => python/examples}/demo.py | 0 pyproject.toml => python/pyproject.toml | 0 .../scripts}/convert_adapters_legacy.py | 0 {scripts => python/scripts}/e2e_smoke.py | 0 {scripts => python/scripts}/fetch_models.py | 0 .../scripts}/generate_example.py | 0 .../scripts}/strip_vision_weights.py | 0 {scripts => python/scripts}/upload_hf.py | 0 .../scripts}/verify_alignment.py | 0 {src => python/src}/edge0/__init__.py | 0 {src => python/src}/edge0/__main__.py | 0 .../src}/edge0/adapters/__init__.py | 0 {src => python/src}/edge0/adapters/lora.py | 0 {src => python/src}/edge0/attention/spec.py | 0 .../src}/edge0/backends/__init__.py | 0 {src => python/src}/edge0/backends/base.py | 0 .../src}/edge0/backends/mlx/__init__.py | 0 .../src}/edge0/backends/mlx/_impl/__init__.py | 0 .../backends/mlx/_impl/bailing_hybrid.py | 0 .../src}/edge0/backends/mlx/_impl/qwen3_5.py | 0 .../edge0/backends/mlx/_impl/qwen3_5_moe.py | 0 .../edge0/backends/mlx/_impl/qwen3_next.py | 0 .../src}/edge0/backends/mlx/backend.py | 0 {src => python/src}/edge0/backends/mlx/io.py | 0 .../src}/edge0/backends/mlx/quant.py | 0 {src => python/src}/edge0/cli.py | 4 ++-- {src => python/src}/edge0/config.py | 0 {src => python/src}/edge0/engine/__init__.py | 0 {src => python/src}/edge0/engine/base.py | 0 {src => python/src}/edge0/engine/hooks.py | 0 {src => python/src}/edge0/engine/ling.py | 0 {src => python/src}/edge0/engine/qwen.py | 0 {src => python/src}/edge0/models/__init__.py | 0 {src => python/src}/edge0/models/base.py | 2 +- .../src}/edge0/models/edge0_35b/__init__.py | 0 .../src}/edge0/models/edge0_8b/__init__.py | 0 {src => python/src}/edge0/moe/__init__.py | 0 {src => python/src}/edge0/moe/routing.py | 0 {src => python/src}/edge0/moe/spec.py | 0 .../src}/edge0/prerouter/__init__.py | 0 {src => python/src}/edge0/prerouter/heads.py | 0 .../src}/edge0/prerouter/install.py | 0 {src => python/src}/edge0/prerouter/spec.py | 0 {src => python/src}/edge0/prerouter/stager.py | 0 {src => python/src}/edge0/prerouter/state.py | 0 {src => python/src}/edge0/registry.py | 0 {src => python/src}/edge0/sampling.py | 0 {src => python/src}/edge0/server/__init__.py | 0 {src => python/src}/edge0/server/app.py | 0 {src => python/src}/edge0/server/chat.py | 0 .../src}/edge0/streaming/__init__.py | 0 {src => python/src}/edge0/streaming/cache.py | 0 .../src}/edge0/streaming/install.py | 0 {src => python/src}/edge0/streaming/layer.py | 0 {src => python/src}/edge0/streaming/mmap.py | 0 .../src}/edge0/streaming/options.py | 0 {tests => python/tests}/test_cli_flags.py | 0 {tests => python/tests}/test_e2e_slow.py | 0 {tests => python/tests}/test_moe_spec.py | 0 {tests => python/tests}/test_prefill_hook.py | 0 {tests => python/tests}/test_registry.py | 0 {tests => python/tests}/test_repo_hygiene.py | 8 ++++--- {tests => python/tests}/test_sampling.py | 0 {tests => python/tests}/test_server.py | 0 .../tests}/test_streaming_math.py | 0 74 files changed, 111 insertions(+), 9 deletions(-) create mode 100644 android/README.md create mode 100644 ios/README.md create mode 100644 macos/README.md create mode 100644 python/README.md rename {examples => python/examples}/bench.py (100%) rename {examples => python/examples}/demo.py (100%) rename pyproject.toml => python/pyproject.toml (100%) rename {scripts => python/scripts}/convert_adapters_legacy.py (100%) rename {scripts => python/scripts}/e2e_smoke.py (100%) rename {scripts => python/scripts}/fetch_models.py (100%) rename {scripts => python/scripts}/generate_example.py (100%) rename {scripts => python/scripts}/strip_vision_weights.py (100%) rename {scripts => python/scripts}/upload_hf.py (100%) rename {scripts => python/scripts}/verify_alignment.py (100%) rename {src => python/src}/edge0/__init__.py (100%) rename {src => python/src}/edge0/__main__.py (100%) rename {src => python/src}/edge0/adapters/__init__.py (100%) rename {src => python/src}/edge0/adapters/lora.py (100%) rename {src => python/src}/edge0/attention/spec.py (100%) rename {src => python/src}/edge0/backends/__init__.py (100%) rename {src => python/src}/edge0/backends/base.py (100%) rename {src => python/src}/edge0/backends/mlx/__init__.py (100%) rename {src => python/src}/edge0/backends/mlx/_impl/__init__.py (100%) rename {src => python/src}/edge0/backends/mlx/_impl/bailing_hybrid.py (100%) rename {src => python/src}/edge0/backends/mlx/_impl/qwen3_5.py (100%) rename {src => python/src}/edge0/backends/mlx/_impl/qwen3_5_moe.py (100%) rename {src => python/src}/edge0/backends/mlx/_impl/qwen3_next.py (100%) rename {src => python/src}/edge0/backends/mlx/backend.py (100%) rename {src => python/src}/edge0/backends/mlx/io.py (100%) rename {src => python/src}/edge0/backends/mlx/quant.py (100%) rename {src => python/src}/edge0/cli.py (99%) rename {src => python/src}/edge0/config.py (100%) rename {src => python/src}/edge0/engine/__init__.py (100%) rename {src => python/src}/edge0/engine/base.py (100%) rename {src => python/src}/edge0/engine/hooks.py (100%) rename {src => python/src}/edge0/engine/ling.py (100%) rename {src => python/src}/edge0/engine/qwen.py (100%) rename {src => python/src}/edge0/models/__init__.py (100%) rename {src => python/src}/edge0/models/base.py (99%) rename {src => python/src}/edge0/models/edge0_35b/__init__.py (100%) rename {src => python/src}/edge0/models/edge0_8b/__init__.py (100%) rename {src => python/src}/edge0/moe/__init__.py (100%) rename {src => python/src}/edge0/moe/routing.py (100%) rename {src => python/src}/edge0/moe/spec.py (100%) rename {src => python/src}/edge0/prerouter/__init__.py (100%) rename {src => python/src}/edge0/prerouter/heads.py (100%) rename {src => python/src}/edge0/prerouter/install.py (100%) rename {src => python/src}/edge0/prerouter/spec.py (100%) rename {src => python/src}/edge0/prerouter/stager.py (100%) rename {src => python/src}/edge0/prerouter/state.py (100%) rename {src => python/src}/edge0/registry.py (100%) rename {src => python/src}/edge0/sampling.py (100%) rename {src => python/src}/edge0/server/__init__.py (100%) rename {src => python/src}/edge0/server/app.py (100%) rename {src => python/src}/edge0/server/chat.py (100%) rename {src => python/src}/edge0/streaming/__init__.py (100%) rename {src => python/src}/edge0/streaming/cache.py (100%) rename {src => python/src}/edge0/streaming/install.py (100%) rename {src => python/src}/edge0/streaming/layer.py (100%) rename {src => python/src}/edge0/streaming/mmap.py (100%) rename {src => python/src}/edge0/streaming/options.py (100%) rename {tests => python/tests}/test_cli_flags.py (100%) rename {tests => python/tests}/test_e2e_slow.py (100%) rename {tests => python/tests}/test_moe_spec.py (100%) rename {tests => python/tests}/test_prefill_hook.py (100%) rename {tests => python/tests}/test_registry.py (100%) rename {tests => python/tests}/test_repo_hygiene.py (91%) rename {tests => python/tests}/test_sampling.py (100%) rename {tests => python/tests}/test_server.py (100%) rename {tests => python/tests}/test_streaming_math.py (100%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9c3a3ad..78cfea2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,6 +9,9 @@ jobs: test-macos: name: Unit tests (macOS + MLX) runs-on: macos-14 + defaults: + run: + working-directory: python steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -22,6 +25,9 @@ jobs: hygiene: name: Repo hygiene (pure Python, no MLX) runs-on: ubuntu-latest + defaults: + run: + working-directory: python steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 diff --git a/.gitignore b/.gitignore index c3b7432..b41702c 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,4 @@ artifacts/ # Bench / runtime logs *.log -scripts/alignment_results.json +python/scripts/alignment_results.json diff --git a/android/README.md b/android/README.md new file mode 100644 index 0000000..6068ea4 --- /dev/null +++ b/android/README.md @@ -0,0 +1,23 @@ +# edge0 — Android runtime + +> ✅ Open-sourced on **2026-09-30**. This README will be updated with +> full build & usage docs together with the source drop. + +The Android runtime runs the edge0 models fully on-device on Android +(Kotlin app + native inference engine): local streaming MoE inference +with the same recipe as the rest of the framework — SSD expert offload, +Recover-LoRA, and prerouter routing prediction. + +The next milestone is the **unified inference framework** — one access +layer, runtime auto-adapting to the hardware platform (iOS / macOS / +Android / Windows / Python) — targeted for the **end of October 2026**; +see the [roadmap](../README.md#-roadmap). + +See the [repository README](../README.md) for the multi-platform picture. + +--- + +> ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 +> 与使用文档。Android runtime：Kotlin App + 原生推理引擎的端侧推理。 +> **统一推理框架**将于 **2026 年 10 月底**发布，详见 +> [根 README 路线图](../README_zh.md)。 diff --git a/docs/adding-a-model.md b/docs/adding-a-model.md index 6af69b4..f14ac8c 100644 --- a/docs/adding-a-model.md +++ b/docs/adding-a-model.md @@ -42,7 +42,7 @@ The subclass only needs to implement the class method `_defaults(model_dir) -> C Key points: -- The repository root ships `scripts/convert_adapters_legacy.py`, which one-shot converts the training npz exports into safetensors artifacts under `artifacts/`; `ModelConfig.artifact(name)` (`models/base.py:28–30`) returns the absolute paths of those artifacts, and the adapter uses them to fill in the LoRA / prerouter weight paths. +- The Python project root (`python/`) ships `scripts/convert_adapters_legacy.py`, which one-shot converts the training npz exports into safetensors artifacts under `artifacts/`; `ModelConfig.artifact(name)` (`models/base.py:28–30`) returns the absolute paths of those artifacts, and the adapter uses them to fill in the LoRA / prerouter weight paths. - LoRA overrides go through `resolve_lora` (`models/base.py:88–95`): a bare model name resolves to that tier's artifacts, `"model_dir"` means "keep the training weights in place", and an empty string disables it. ## Step 3: `build_model` / `build_engine` diff --git a/docs/prerouter.md b/docs/prerouter.md index 460a1da..8744fa4 100644 --- a/docs/prerouter.md +++ b/docs/prerouter.md @@ -73,7 +73,7 @@ per `start_layer`/`owners`. Family differences are expressed through three hooks | weight file | `prerouter_edge0_35b.safetensors` | `prerouter_edge0_8b.safetensors` | Both weight files are resolved from the model directory by default, with -`artifacts/` (repo root, gitignored) as the fallback — see the README's +`artifacts/` (Python project root, gitignored) as the fallback — see the README's "Models and adapters" section. ## Tests and regression diff --git a/ios/README.md b/ios/README.md new file mode 100644 index 0000000..389fa70 --- /dev/null +++ b/ios/README.md @@ -0,0 +1,23 @@ +# edge0 — iOS runtime + +> ✅ Open-sourced on **2026-09-30**. This README will be updated with +> full build & usage docs together with the source drop. + +The iOS runtime runs the edge0 models fully on-device on iPhone +(Swift / SwiftUI + MLX Swift): local streaming MoE inference with the +same recipe as the rest of the framework — SSD expert offload, +Recover-LoRA, and prerouter routing prediction. + +The next milestone is the **unified inference framework** — one access +layer, runtime auto-adapting to the hardware platform (iOS / macOS / +Android / Windows / Python) — targeted for the **end of October 2026**; +see the [roadmap](../README.md#-roadmap). + +See the [repository README](../README.md) for the multi-platform picture. + +--- + +> ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 +> 与使用文档。iOS runtime：Swift / SwiftUI + MLX Swift 的 iPhone 端侧 +> 推理。**统一推理框架**将于 **2026 年 10 月底**发布，详见 +> [根 README 路线图](../README_zh.md)。 diff --git a/macos/README.md b/macos/README.md new file mode 100644 index 0000000..35e7ddd --- /dev/null +++ b/macos/README.md @@ -0,0 +1,24 @@ +# edge0 — macOS runtime + +> ✅ Open-sourced on **2026-09-30**. This README will be updated with +> full build & usage docs together with the source drop. + +The macOS runtime is the local CLI / daemon / desktop app for edge0 on +Apple Silicon (Rust + Tauri): an Ollama-style UX that pulls the edge0 +models ([Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview), +[Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)) +and serves them locally behind an OpenAI-compatible API. + +The next milestone is the **unified inference framework** — one access +layer, runtime auto-adapting to the hardware platform (iOS / macOS / +Android / Windows / Python) — targeted for the **end of October 2026**; +see the [roadmap](../README.md#-roadmap). + +See the [repository README](../README.md) for the multi-platform picture. + +--- + +> ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 +> 与使用文档。macOS runtime：Rust + Tauri 的本地 CLI / daemon / 桌面 +> App，Ollama 式体验 + OpenAI 兼容 API。**统一推理框架**将于 **2026 年 +> 10 月底**发布，详见[根 README 路线图](../README_zh.md)。 diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000..2bee612 --- /dev/null +++ b/python/README.md @@ -0,0 +1,24 @@ +# edge0 — Python framework + +The Python implementation of the edge0 streaming MoE inference framework +(SSD expert offload + Recover-LoRA + prerouter routing prediction), +shipping the MLX backend for Apple Silicon. + +This directory is one platform of the edge0 multi-platform repo — see the +[repository README](../README.md) for the full picture (macOS / iOS / +Android runtimes and the unified-framework roadmap). + +## Quick start + +```bash +python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]' +edge0 demo edge0-8b # after downloading a model; see ../README.md +``` + +- Docs: [`docs/`](../docs/) at the repo root (architecture, attention, MoE, streaming, prerouter) +- Examples: [`examples/`](examples/) · Scripts: [`scripts/`](scripts/) +- Tests: `pytest` (unit) · `pytest -m slow` (real weights) + +## License + +Apache-2.0, including vendored third-party code (see [NOTICE](../NOTICE)). diff --git a/examples/bench.py b/python/examples/bench.py similarity index 100% rename from examples/bench.py rename to python/examples/bench.py diff --git a/examples/demo.py b/python/examples/demo.py similarity index 100% rename from examples/demo.py rename to python/examples/demo.py diff --git a/pyproject.toml b/python/pyproject.toml similarity index 100% rename from pyproject.toml rename to python/pyproject.toml diff --git a/scripts/convert_adapters_legacy.py b/python/scripts/convert_adapters_legacy.py similarity index 100% rename from scripts/convert_adapters_legacy.py rename to python/scripts/convert_adapters_legacy.py diff --git a/scripts/e2e_smoke.py b/python/scripts/e2e_smoke.py similarity index 100% rename from scripts/e2e_smoke.py rename to python/scripts/e2e_smoke.py diff --git a/scripts/fetch_models.py b/python/scripts/fetch_models.py similarity index 100% rename from scripts/fetch_models.py rename to python/scripts/fetch_models.py diff --git a/scripts/generate_example.py b/python/scripts/generate_example.py similarity index 100% rename from scripts/generate_example.py rename to python/scripts/generate_example.py diff --git a/scripts/strip_vision_weights.py b/python/scripts/strip_vision_weights.py similarity index 100% rename from scripts/strip_vision_weights.py rename to python/scripts/strip_vision_weights.py diff --git a/scripts/upload_hf.py b/python/scripts/upload_hf.py similarity index 100% rename from scripts/upload_hf.py rename to python/scripts/upload_hf.py diff --git a/scripts/verify_alignment.py b/python/scripts/verify_alignment.py similarity index 100% rename from scripts/verify_alignment.py rename to python/scripts/verify_alignment.py diff --git a/src/edge0/__init__.py b/python/src/edge0/__init__.py similarity index 100% rename from src/edge0/__init__.py rename to python/src/edge0/__init__.py diff --git a/src/edge0/__main__.py b/python/src/edge0/__main__.py similarity index 100% rename from src/edge0/__main__.py rename to python/src/edge0/__main__.py diff --git a/src/edge0/adapters/__init__.py b/python/src/edge0/adapters/__init__.py similarity index 100% rename from src/edge0/adapters/__init__.py rename to python/src/edge0/adapters/__init__.py diff --git a/src/edge0/adapters/lora.py b/python/src/edge0/adapters/lora.py similarity index 100% rename from src/edge0/adapters/lora.py rename to python/src/edge0/adapters/lora.py diff --git a/src/edge0/attention/spec.py b/python/src/edge0/attention/spec.py similarity index 100% rename from src/edge0/attention/spec.py rename to python/src/edge0/attention/spec.py diff --git a/src/edge0/backends/__init__.py b/python/src/edge0/backends/__init__.py similarity index 100% rename from src/edge0/backends/__init__.py rename to python/src/edge0/backends/__init__.py diff --git a/src/edge0/backends/base.py b/python/src/edge0/backends/base.py similarity index 100% rename from src/edge0/backends/base.py rename to python/src/edge0/backends/base.py diff --git a/src/edge0/backends/mlx/__init__.py b/python/src/edge0/backends/mlx/__init__.py similarity index 100% rename from src/edge0/backends/mlx/__init__.py rename to python/src/edge0/backends/mlx/__init__.py diff --git a/src/edge0/backends/mlx/_impl/__init__.py b/python/src/edge0/backends/mlx/_impl/__init__.py similarity index 100% rename from src/edge0/backends/mlx/_impl/__init__.py rename to python/src/edge0/backends/mlx/_impl/__init__.py diff --git a/src/edge0/backends/mlx/_impl/bailing_hybrid.py b/python/src/edge0/backends/mlx/_impl/bailing_hybrid.py similarity index 100% rename from src/edge0/backends/mlx/_impl/bailing_hybrid.py rename to python/src/edge0/backends/mlx/_impl/bailing_hybrid.py diff --git a/src/edge0/backends/mlx/_impl/qwen3_5.py b/python/src/edge0/backends/mlx/_impl/qwen3_5.py similarity index 100% rename from src/edge0/backends/mlx/_impl/qwen3_5.py rename to python/src/edge0/backends/mlx/_impl/qwen3_5.py diff --git a/src/edge0/backends/mlx/_impl/qwen3_5_moe.py b/python/src/edge0/backends/mlx/_impl/qwen3_5_moe.py similarity index 100% rename from src/edge0/backends/mlx/_impl/qwen3_5_moe.py rename to python/src/edge0/backends/mlx/_impl/qwen3_5_moe.py diff --git a/src/edge0/backends/mlx/_impl/qwen3_next.py b/python/src/edge0/backends/mlx/_impl/qwen3_next.py similarity index 100% rename from src/edge0/backends/mlx/_impl/qwen3_next.py rename to python/src/edge0/backends/mlx/_impl/qwen3_next.py diff --git a/src/edge0/backends/mlx/backend.py b/python/src/edge0/backends/mlx/backend.py similarity index 100% rename from src/edge0/backends/mlx/backend.py rename to python/src/edge0/backends/mlx/backend.py diff --git a/src/edge0/backends/mlx/io.py b/python/src/edge0/backends/mlx/io.py similarity index 100% rename from src/edge0/backends/mlx/io.py rename to python/src/edge0/backends/mlx/io.py diff --git a/src/edge0/backends/mlx/quant.py b/python/src/edge0/backends/mlx/quant.py similarity index 100% rename from src/edge0/backends/mlx/quant.py rename to python/src/edge0/backends/mlx/quant.py diff --git a/src/edge0/cli.py b/python/src/edge0/cli.py similarity index 99% rename from src/edge0/cli.py rename to python/src/edge0/cli.py index 97555d8..d2cbb73 100644 --- a/src/edge0/cli.py +++ b/python/src/edge0/cli.py @@ -211,9 +211,9 @@ def cmd_serve(args) -> int: def cmd_convert(args) -> int: import runpy sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve() - .parents[1] / "scripts")) + .parents[2] / "scripts")) runpy.run_path( - str(__import__("pathlib").Path(__file__).resolve().parents[1] + str(__import__("pathlib").Path(__file__).resolve().parents[2] / "scripts" / "convert_adapters_legacy.py"), run_name="__main__", ) diff --git a/src/edge0/config.py b/python/src/edge0/config.py similarity index 100% rename from src/edge0/config.py rename to python/src/edge0/config.py diff --git a/src/edge0/engine/__init__.py b/python/src/edge0/engine/__init__.py similarity index 100% rename from src/edge0/engine/__init__.py rename to python/src/edge0/engine/__init__.py diff --git a/src/edge0/engine/base.py b/python/src/edge0/engine/base.py similarity index 100% rename from src/edge0/engine/base.py rename to python/src/edge0/engine/base.py diff --git a/src/edge0/engine/hooks.py b/python/src/edge0/engine/hooks.py similarity index 100% rename from src/edge0/engine/hooks.py rename to python/src/edge0/engine/hooks.py diff --git a/src/edge0/engine/ling.py b/python/src/edge0/engine/ling.py similarity index 100% rename from src/edge0/engine/ling.py rename to python/src/edge0/engine/ling.py diff --git a/src/edge0/engine/qwen.py b/python/src/edge0/engine/qwen.py similarity index 100% rename from src/edge0/engine/qwen.py rename to python/src/edge0/engine/qwen.py diff --git a/src/edge0/models/__init__.py b/python/src/edge0/models/__init__.py similarity index 100% rename from src/edge0/models/__init__.py rename to python/src/edge0/models/__init__.py diff --git a/src/edge0/models/base.py b/python/src/edge0/models/base.py similarity index 99% rename from src/edge0/models/base.py rename to python/src/edge0/models/base.py index d3dead5..bb1d73d 100644 --- a/src/edge0/models/base.py +++ b/python/src/edge0/models/base.py @@ -19,7 +19,7 @@ from edge0.prerouter.spec import PrerouterSpec from edge0.streaming.options import LayerOptions -# Repository-root artifacts (gitignored; produced once by +# Project-root artifacts (gitignored; produced once by # scripts/convert_adapters_legacy.py from the training npz exports, then # the source npz are discarded). ARTIFACTS_DIR = Path(__file__).resolve().parents[3] / "artifacts" diff --git a/src/edge0/models/edge0_35b/__init__.py b/python/src/edge0/models/edge0_35b/__init__.py similarity index 100% rename from src/edge0/models/edge0_35b/__init__.py rename to python/src/edge0/models/edge0_35b/__init__.py diff --git a/src/edge0/models/edge0_8b/__init__.py b/python/src/edge0/models/edge0_8b/__init__.py similarity index 100% rename from src/edge0/models/edge0_8b/__init__.py rename to python/src/edge0/models/edge0_8b/__init__.py diff --git a/src/edge0/moe/__init__.py b/python/src/edge0/moe/__init__.py similarity index 100% rename from src/edge0/moe/__init__.py rename to python/src/edge0/moe/__init__.py diff --git a/src/edge0/moe/routing.py b/python/src/edge0/moe/routing.py similarity index 100% rename from src/edge0/moe/routing.py rename to python/src/edge0/moe/routing.py diff --git a/src/edge0/moe/spec.py b/python/src/edge0/moe/spec.py similarity index 100% rename from src/edge0/moe/spec.py rename to python/src/edge0/moe/spec.py diff --git a/src/edge0/prerouter/__init__.py b/python/src/edge0/prerouter/__init__.py similarity index 100% rename from src/edge0/prerouter/__init__.py rename to python/src/edge0/prerouter/__init__.py diff --git a/src/edge0/prerouter/heads.py b/python/src/edge0/prerouter/heads.py similarity index 100% rename from src/edge0/prerouter/heads.py rename to python/src/edge0/prerouter/heads.py diff --git a/src/edge0/prerouter/install.py b/python/src/edge0/prerouter/install.py similarity index 100% rename from src/edge0/prerouter/install.py rename to python/src/edge0/prerouter/install.py diff --git a/src/edge0/prerouter/spec.py b/python/src/edge0/prerouter/spec.py similarity index 100% rename from src/edge0/prerouter/spec.py rename to python/src/edge0/prerouter/spec.py diff --git a/src/edge0/prerouter/stager.py b/python/src/edge0/prerouter/stager.py similarity index 100% rename from src/edge0/prerouter/stager.py rename to python/src/edge0/prerouter/stager.py diff --git a/src/edge0/prerouter/state.py b/python/src/edge0/prerouter/state.py similarity index 100% rename from src/edge0/prerouter/state.py rename to python/src/edge0/prerouter/state.py diff --git a/src/edge0/registry.py b/python/src/edge0/registry.py similarity index 100% rename from src/edge0/registry.py rename to python/src/edge0/registry.py diff --git a/src/edge0/sampling.py b/python/src/edge0/sampling.py similarity index 100% rename from src/edge0/sampling.py rename to python/src/edge0/sampling.py diff --git a/src/edge0/server/__init__.py b/python/src/edge0/server/__init__.py similarity index 100% rename from src/edge0/server/__init__.py rename to python/src/edge0/server/__init__.py diff --git a/src/edge0/server/app.py b/python/src/edge0/server/app.py similarity index 100% rename from src/edge0/server/app.py rename to python/src/edge0/server/app.py diff --git a/src/edge0/server/chat.py b/python/src/edge0/server/chat.py similarity index 100% rename from src/edge0/server/chat.py rename to python/src/edge0/server/chat.py diff --git a/src/edge0/streaming/__init__.py b/python/src/edge0/streaming/__init__.py similarity index 100% rename from src/edge0/streaming/__init__.py rename to python/src/edge0/streaming/__init__.py diff --git a/src/edge0/streaming/cache.py b/python/src/edge0/streaming/cache.py similarity index 100% rename from src/edge0/streaming/cache.py rename to python/src/edge0/streaming/cache.py diff --git a/src/edge0/streaming/install.py b/python/src/edge0/streaming/install.py similarity index 100% rename from src/edge0/streaming/install.py rename to python/src/edge0/streaming/install.py diff --git a/src/edge0/streaming/layer.py b/python/src/edge0/streaming/layer.py similarity index 100% rename from src/edge0/streaming/layer.py rename to python/src/edge0/streaming/layer.py diff --git a/src/edge0/streaming/mmap.py b/python/src/edge0/streaming/mmap.py similarity index 100% rename from src/edge0/streaming/mmap.py rename to python/src/edge0/streaming/mmap.py diff --git a/src/edge0/streaming/options.py b/python/src/edge0/streaming/options.py similarity index 100% rename from src/edge0/streaming/options.py rename to python/src/edge0/streaming/options.py diff --git a/tests/test_cli_flags.py b/python/tests/test_cli_flags.py similarity index 100% rename from tests/test_cli_flags.py rename to python/tests/test_cli_flags.py diff --git a/tests/test_e2e_slow.py b/python/tests/test_e2e_slow.py similarity index 100% rename from tests/test_e2e_slow.py rename to python/tests/test_e2e_slow.py diff --git a/tests/test_moe_spec.py b/python/tests/test_moe_spec.py similarity index 100% rename from tests/test_moe_spec.py rename to python/tests/test_moe_spec.py diff --git a/tests/test_prefill_hook.py b/python/tests/test_prefill_hook.py similarity index 100% rename from tests/test_prefill_hook.py rename to python/tests/test_prefill_hook.py diff --git a/tests/test_registry.py b/python/tests/test_registry.py similarity index 100% rename from tests/test_registry.py rename to python/tests/test_registry.py diff --git a/tests/test_repo_hygiene.py b/python/tests/test_repo_hygiene.py similarity index 91% rename from tests/test_repo_hygiene.py rename to python/tests/test_repo_hygiene.py index c8ef1d0..a6ed8ef 100644 --- a/tests/test_repo_hygiene.py +++ b/python/tests/test_repo_hygiene.py @@ -16,13 +16,15 @@ import pathlib import re -ROOT = pathlib.Path(__file__).resolve().parents[1] +# Repo root (this file lives at /python/tests/): the hygiene guards +# scan the whole multi-platform repo, not just the Python subproject. +ROOT = pathlib.Path(__file__).resolve().parents[2] _EXCLUDE_DIRS = {".git", ".venv", ".pytest_cache", "__pycache__", ".agents", ".codex"} _EXCLUDE_PARTS = {"egg-info"} _EXCLUDE_SUFFIXES = {".pyc", ".so", ".dylib", ".bin", ".safetensors", ".npz", - ".png", ".jpg", ".jpeg", ".gif", ".ico"} + ".png", ".jpg", ".jpeg", ".gif", ".ico", ".pdf"} # NOTE: patterns are assembled with + so this file never contains the # literal string it is looking for. @@ -85,7 +87,7 @@ def test_no_hardcoded_local_paths(): def test_mlx_imports_stay_inside_the_backend(): bad = [] - for path in (ROOT / "src" / "edge0").rglob("*.py"): + for path in (ROOT / "python" / "src" / "edge0").rglob("*.py"): rel = path.relative_to(ROOT) text = path.read_text(encoding="utf-8") if _MLX_IMPORT.search(text) and "backends/mlx" not in rel.as_posix(): diff --git a/tests/test_sampling.py b/python/tests/test_sampling.py similarity index 100% rename from tests/test_sampling.py rename to python/tests/test_sampling.py diff --git a/tests/test_server.py b/python/tests/test_server.py similarity index 100% rename from tests/test_server.py rename to python/tests/test_server.py diff --git a/tests/test_streaming_math.py b/python/tests/test_streaming_math.py similarity index 100% rename from tests/test_streaming_math.py rename to python/tests/test_streaming_math.py From cafd81be4e2403d5d195eff4c1daf9b03ae04ecc Mon Sep 17 00:00:00 2001 From: Bin Date: Mon, 28 Sep 2026 16:37:02 +0800 Subject: [PATCH 02/13] docs(readme): restructure EN/ZH READMEs in sglang/vllm style Reorganize both READMEs as News -> About -> Getting Started -> Roadmap -> Contributing -> Citation -> Contact Us: - News: 2026-09-30 three-platform engine release (iOS / macOS / Android, source open-sourced in-repo), --prefill-ondemand, the arXiv report, and the initial release - About: core mechanisms, a platform matrix (Python available now, three engines open-sourced, Windows on the roadmap), models, design (backend isolation scoped to the Python framework), quality and benchmark tables carried over verbatim - Getting Started: full Python walkthrough on the new python/ paths; macOS / iOS / Android point at each directory's README - Roadmap: the unified inference framework (one access layer, runtime auto-adapting to iOS / macOS / Android / Windows / Python) by the end of October 2026, plus CUDA backend and Windows - Citation: BibTeX for arXiv 2609.18063 - Contact Us: placeholder channels (email / Discord / WeChat) with a TODO for the team to fill in Co-Authored-By: Claude Code --- README.md | 291 +++++++++++++++++++++++++++++++++++---------------- README_zh.md | 272 +++++++++++++++++++++++++++++++++-------------- 2 files changed, 393 insertions(+), 170 deletions(-) diff --git a/README.md b/README.md index 5e0240a..110e96c 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,8 @@ **An open-source streaming MoE inference framework — SSD expert offload + Recover-LoRA + prerouter routing prediction.** +**Python** · **macOS** · **iOS** · **Android** — one recipe, every device. + [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) [![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) @@ -18,12 +20,48 @@ English | [中文](README_zh.md) +## 📰 News + +- **[2026-09-30]** We released the **edge0 inference engines for three platforms — iOS, macOS and Android** — so users get the best inference experience across architectures and platforms. The source is open-sourced in this repo ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/)) — see each directory's README for details. The **unified inference framework** follows by the **end of October 2026**; see the [Roadmap](#-roadmap). +- **[2026-09-20]** `--prefill-ondemand` lands for `edge0 demo / chat / serve` ([#112](https://github.com/Edge0-AI/Edge0/pull/112)). +- **[2026-09-16]** Our technical report is on arXiv: [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063). +- **[2026-09-08]** Initial open-source release of **edge0**, together with both model tiers — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) and [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — on Hugging Face and ModelScope. + +## 🤔 About + **edge0** is an open-source streaming MoE inference framework. It generalizes the production-proven recipe — **SSD expert offload + Recover-LoRA + prerouter routing prediction** — into an extensible -framework. The backend is isolated by design: the current MLX backend -runs on Apple Silicon, and additional platforms (CUDA, …) plug into the -same core abstractions. +framework that runs large sparse-MoE models on consumer hardware: peak +memory is bounded by the *active* expert set, not the parameter count. + +### Core mechanisms + +- **SSD expert offload**: expert weights are streamed from storage on + demand; peak memory is bounded by the active set, not the parameter + count. +- **Prerouter**: a trained head predicts expert routing one step + ahead, so expert loads overlap the forward pass instead of stalling + it — **up to +59%** decode throughput; the gain grows with storage + latency, model size, and routed width *K*. +- **Recover-LoRA**: the int4 base is frozen and LoRA adapters are + trained by distillation from the FP teacher, recovering most of the + quantization loss at 4-bit (see [Quality](#quality)). Adapters stay + unmerged: one read-only base serves multiple adapter sets. + +### Platforms + +One repo, one recipe, per-platform runtimes: + +| Platform | Directory | Stack | Status | +|---|---|---|---| +| **Python** (macOS · Apple Silicon) | [`python/`](python/) | Python + MLX | ✅ Available now | +| **macOS** app & CLI | [`macos/`](macos/) | Rust | ✅ Open-sourced (2026-09-30) | +| **iOS** app | [`ios/`](ios/) | Swift + MLX Swift | ✅ Open-sourced (2026-09-30) | +| **Android** app & engine | [`android/`](android/) | Kotlin + native engine | ✅ Open-sourced (2026-09-30) | +| **Windows** | — | — | 🗺️ On the roadmap | + +### Models Two model tiers ship with the framework. Each tier is an end-to-end release: the released checkpoint, the trained LoRA adapters, and the @@ -40,32 +78,19 @@ LoRA and prerouter training done for this framework — the adapter files are co-located with each checkpoint and load automatically, so `edge0 serve ` runs the trained pipeline out of the box. -## Requirements - -- **OS / hardware**: the MLX backend runs on macOS with Apple Silicon - (M1/M2/M3/M4). The CUDA backend is on the roadmap — no other - platforms are supported yet. -- **Python**: 3.10+ (3.12 recommended). -- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see - `pyproject.toml`). Garbled, mixed-language output on Apple A18 / A18 Pro - means an older `mlx`: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'` - ([#8](https://github.com/Edge0-AI/Edge0/issues/8)). -- **Memory**: ~2.9 GB peak active memory for `edge0-35b`, ~1.0 GB for - `edge0-8b` (short contexts; see [Benchmark](#benchmark)). Add - headroom for the OS, tokenizer, and long-context KV growth. -- **Disk**: the 4-bit checkpoints are ~23 GB (`edge0-35b`) and ~4.2 GB - (`edge0-8b`); expert weights are mmapped and read on demand, they are - not loaded into RAM up front. - -## Design +### Design - **transformers-style usage**: `AutoModel` / `AutoConfig` / `AutoEngine` resolve the tier from the model name; -- **Backend isolation**: all MLX code lives under `edge0/backends/mlx/`; - the core logic (model specs, prerouter, streaming expert pool, server) - depends only on the backend facade (`edge0/backends/base.py`), so a new - backend implements the same facade (`backends/cuda/` is a reserved - slot) with zero changes to core code; +- **Backend isolation (Python framework)**: within the Python framework, + all MLX code lives under `python/src/edge0/backends/mlx/`; the core + logic (model specs, prerouter, streaming expert pool, server) depends + only on the backend facade (`backends/base.py`), so a new backend + implements the same facade (`backends/cuda/` is a reserved slot) with + zero changes to core code. The iOS / macOS / Android engines ship + platform-native stacks today — bringing every platform under one + access layer is exactly what the unified inference framework + (see [Roadmap](#-roadmap)) will deliver; - **Adapters as safetensors**: LoRA and prerouter weights are `.safetensors` files with provenance metadata (source, version, owner layers), resolved from the model directory or `artifacts/`; @@ -74,34 +99,79 @@ are co-located with each checkpoint and load automatically, so and that model's adapters; upgrading adapters swaps adapter files only — the base stays read-only and is never merged. -## Core mechanisms +### Quality -- **SSD expert offload**: expert weights are streamed from storage on - demand; peak memory is bounded by the active set, not the parameter - count. -- **Prerouter**: a trained head predicts expert routing one step - ahead, so expert loads overlap the forward pass instead of stalling - it — **up to +59%** decode throughput; the gain grows with storage - latency, model size, and routed width *K*. -- **Recover-LoRA**: the int4 base is frozen and LoRA adapters are - trained by distillation from the FP teacher, recovering most of the - quantization loss at 4-bit (see [Quality](#quality)). Adapters stay - unmerged: one read-only base serves multiple adapter sets. +All benchmarks were run by us with [OpenCompass](https://github.com/open-compass/opencompass) +under identical settings and parameters for both the edge0 models (int4 + +trained adapters + prerouter routing) and the original fp16 base models. +The loss of the edge0 pipeline is small: **3.9 points on average for +edge0-35b, 2.8 for edge0-8b** (MMLU-Pro is even above the base). Max 100: + +| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) | +|---|---:|---:|---:|---:| +| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 | +| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 | +| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 | +| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 | +| IFBench | 57.9 | 61.7 | 53.9 | 60.6 | +| **Average** | **79.2** | **83.2** | **69.9** | **72.7** | + +### Benchmark + +Measured with `python/examples/bench.py` (3.3k-token prompt prefill → 10 +sampled warmup steps → 200 timed sampled decode tokens, 2 runs per tier): + +| Tier | Decode speed | Prefill throughput (cold / warm)* | Peak active memory | Test machine | +|---|---|---|---|---| +| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB | +| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB | + +*Cold = first request after process start (expert weights fault in from +SSD); warm = subsequent requests (page cache resident). Prefill numbers +are throughput over a ~3.3k-token prompt (`BENCH_LONG=1`).* -## Quick start +Reproduce: -### 1) Install +```bash +cd python +python examples/bench.py edge0-35b # via $EDGE0_35B_MODEL +python examples/bench.py edge0-8b # via $EDGE0_8B_MODEL +``` + +## 🚀 Getting Started + +### Python (macOS · Apple Silicon) + +#### Requirements + +- **OS / hardware**: the MLX backend runs on macOS with Apple Silicon + (M1/M2/M3/M4). The CUDA backend is on the roadmap — no other + platforms are supported by the Python framework yet. +- **Python**: 3.10+ (3.12 recommended). +- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see + `python/pyproject.toml`). Garbled, mixed-language output on Apple A18 / + A18 Pro means an older `mlx`: `pip install 'mlx==0.30.6' + 'mlx-metal==0.30.6'` ([#8](https://github.com/Edge0-AI/Edge0/issues/8)). +- **Memory**: ~2.9 GB peak active memory for `edge0-35b`, ~1.0 GB for + `edge0-8b` (short contexts; see [Benchmark](#benchmark)). Add + headroom for the OS, tokenizer, and long-context KV growth. +- **Disk**: the 4-bit checkpoints are ~23 GB (`edge0-35b`) and ~4.2 GB + (`edge0-8b`); expert weights are mmapped and read on demand, they are + not loaded into RAM up front. + +#### 1) Install ```bash +cd python # Python >= 3.10; the MLX backend requires macOS with Apple Silicon python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]' ``` -### 2) Download a model +#### 2) Download a model -The two tiers are published on Hugging Face and ModelScope — each repo bundles the -base checkpoint and the trained LoRA + prerouter adapters in **one -directory**, so a single download is a ready-to-run model: +The two tiers are published on Hugging Face and ModelScope — each repo +bundles the base checkpoint and the trained LoRA + prerouter adapters in +**one directory**, so a single download is a ready-to-run model: - [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) (~23 GB) · [ModelScope mirror](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) - [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) (~4.2 GB) · [ModelScope mirror](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview) @@ -125,7 +195,7 @@ models/edge0-35b/ └── prerouter_edge0_35b.safetensors # trained prerouter heads ``` -### 3) Point edge0 at it +#### 3) Point edge0 at it Tier names resolve to local directories via environment variables (where you put the download is up to you): @@ -143,7 +213,7 @@ edge0 demo models/edge0-35b edge0 serve models/edge0-8b ``` -### 4) Run +#### 4) Run ```bash # quick demo @@ -165,7 +235,7 @@ edge0 chat edge0-35b --prompt "Explain streaming inference in one sentence." `python -m edge0 ...` is equivalent to `edge0 ...`. -### Python API +#### Python API ```python from edge0 import AutoEngine @@ -182,10 +252,10 @@ print(engine._tok.decode(tokens)) engine.close() # release mmaps / expert cache ``` -`examples/demo.py` is the same minimal walkthrough (`edge0 demo` runs -this exact path). +`python/examples/demo.py` is the same minimal walkthrough (`edge0 demo` +runs this exact path). -### Models and adapters +#### Models and adapters - **Checkpoint**: the original model directory (`config.json`, `model*.safetensors`, tokenizer). `edge0 serve ` / @@ -195,72 +265,111 @@ this exact path). location automatically: - the model directory (recommended): side by side with the base, e.g. `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`; - - `artifacts/` (repo root, gitignored): convert once from - training-side npz exports via `edge0 convert-adapters --npz-dir ...`. + - `artifacts/` (gitignored): convert once from training-side npz + exports via `edge0 convert-adapters --npz-dir ...`. - The published model repos bundle both the base checkpoint and the current default adapter release, so `scripts/fetch_models.py` produces - a ready-to-run model directory. Check each model's doc page for its + a ready-to-run model directory. Check each model's doc page for its adapter provenance (training data, owner-layer layout). - Both adapters are required for the prerouter + LoRA pipeline; if a file is missing, `edge0` fails with a clear message (or pass `--no-prerouter` / `--no-lora` to run the plain base model). -## Quality -All benchmarks were run by us with [OpenCompass](https://github.com/open-compass/opencompass) -under identical settings and parameters for both the edge0 models (int4 + -trained adapters + prerouter routing) and the original fp16 base models. -The loss of the edge0 pipeline is small: **3.9 points on average for -edge0-35b, 2.8 for edge0-8b** (MMLU-Pro is even above the base). Max 100: +#### Documentation -| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) | -|---|---:|---:|---:|---:| -| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 | -| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 | -| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 | -| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 | -| IFBench | 57.9 | 61.7 | 53.9 | 60.6 | -| **Average** | **79.2** | **83.2** | **69.9** | **72.7** | +- [Architecture](docs/architecture.md) +- [Attention](docs/attention.md) / [MoE](docs/moe.md) / [SSD streaming](docs/streaming.md) / [prerouter](docs/prerouter.md) +- [Adding a model](docs/adding-a-model.md) +- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) +- Technical report: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](paper/main.pdf)) -## Benchmark +### macOS / iOS / Android -Measured with `examples/bench.py` (3.3k-token prompt prefill → 10 sampled -warmup steps → 200 timed sampled decode tokens, 2 runs per tier): +The three platform engines are open-sourced in this repo — more details +in each directory's README: -| Tier | Decode speed | Prefill throughput (cold / warm)* | Peak active memory | Test machine | -|---|---|---|---|---| -| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB | -| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB | +- **macOS**: local CLI / daemon / desktop app (Rust) — see [`macos/README.md`](macos/README.md) +- **iOS**: on-device iPhone app (Swift + MLX Swift) — see [`ios/README.md`](ios/README.md) +- **Android**: on-device app + native engine (Kotlin) — see [`android/README.md`](android/README.md) -*Cold = first request after process start (expert weights fault in from -SSD); warm = subsequent requests (page cache resident). Prefill numbers -are throughput over a ~3.3k-token prompt (`BENCH_LONG=1`).* +The **unified inference framework** — one access layer, runtime +auto-adapting to iOS / macOS / Android / Windows / Python — arrives by +the **end of October 2026**; see the [Roadmap](#-roadmap). -Reproduce: +## 🗺️ Roadmap -```bash -python examples/bench.py edge0-35b # via $EDGE0_35B_MODEL -python examples/bench.py edge0-8b # via $EDGE0_8B_MODEL -``` +- **[End of Oct 2026] edge0 unified inference framework** — we will + open-source a unified inference framework: **one unified access + layer** (a single API across chat / serve / on-device use), with the + **runtime automatically adapting to the hardware platform** — iOS, + macOS, Android, Windows and Python. It builds on the three platform + engines already open-sourced in this repo (`ios/` · `macos/` · + `android/`). +- **Windows** platform support (via the unified framework). +- **CUDA backend** for the Python framework — reserved slot at + `python/src/edge0/backends/cuda/`, core code needs zero changes. +- More model tiers and adapter releases on the existing pipeline. + +## 🤝 Contributing + +Contributions are welcome — issues, PRs, benchmark reports and model +ports all count. -## Tests +**Python framework** (available now): ```bash +cd python +python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]' + pytest # unit tests (no real weights) EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q - # real-weight generation; missing tiers are skipped + # real-weight generation; missing tiers are skipped .venv/bin/python scripts/e2e_smoke.py \ --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b - # staged vs exact consistency + generation smoke + # staged vs exact consistency + generation smoke scripts/generate_example.py # full-pipeline API example -examples/demo.py # minimal API walkthrough +examples/demo.py # minimal API walkthrough ``` -## Documentation +CI runs unit tests (macOS + MLX) and a repo-hygiene suite (no hardcoded +paths, backend-boundary and secret checks) on every PR. -- [Architecture](docs/architecture.md) -- [Attention](docs/attention.md) / [MoE](docs/moe.md) / [SSD streaming](docs/streaming.md) / [prerouter](docs/prerouter.md) -- [Adding a model](docs/adding-a-model.md) -- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) +**Platform runtimes** (macOS / iOS / Android): the directories are +reserved placeholders for now; contribution guides will land together +with the unified framework code. + +Workflow: fork → feature branch → PR against `main`. Please keep the +hygiene suite green and add tests for new behavior. + +## 📖 Citation + +If you find edge0 useful, please cite our technical report: + +```bibtex +@article{lin2026other, + title = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD + with Trained Routing Prediction}, + author = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and + Zeng, Xiaodong}, + journal = {arXiv preprint arXiv:2609.18063}, + year = {2026}, + url = {https://arxiv.org/abs/2609.18063} +} +``` + +## 📧 Contact Us + + + +Community and support channels are coming soon — this section will list +the official ways to reach us: + +- **Email**: _to be added_ +- **Discord**: _to be added_ +- **WeChat group**: _to be added_ + +For bugs and feature requests, please use +[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues). ## License diff --git a/README_zh.md b/README_zh.md index 09f0d5e..6631ef3 100644 --- a/README_zh.md +++ b/README_zh.md @@ -6,6 +6,8 @@ **开源流式 MoE 推理框架 —— SSD 专家 offload + Recover-LoRA + prerouter 路由预判** +**Python** · **macOS** · **iOS** · **Android** —— 一套配方，全端落地。 + [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) [![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) @@ -18,9 +20,53 @@ +## 📰 新闻 + +- **【2026-09-30】** 我们发布了**三端推理引擎**，适配 iOS、macOS、Android + 平台，让用户在不同架构、不同平台上都有最佳的推理体验。三端源码已 + 开源至本仓库（[`ios/`](ios/) · [`macos/`](macos/) · + [`android/`](android/)），更多细节见各目录 README；**统一推理框架** + 将于 **2026 年 10 月底**发布，详见[路线图](#-路线图)。 +- **【2026-09-20】** `edge0 demo / chat / serve` 新增 `--prefill-ondemand` + （[#112](https://github.com/Edge0-AI/Edge0/pull/112)）。 +- **【2026-09-16】** 技术报告上线 arXiv：[The Other Half of the Memory + Wall: Serving 35B MoEs from SSD with Trained Routing + Prediction](https://arxiv.org/abs/2609.18063)。 +- **【2026-09-08】** **edge0** 首次开源发布，两个模型档位 + [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) + 与 [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) + 同步登陆 Hugging Face 与 ModelScope。 + +## 🤔 关于 edge0 + **edge0** 是一个开源的流式 MoE 推理框架：把「SSD 专家 offload + -Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架。后端隔离设计， -当前实现 MLX 后端（Apple Silicon），更多平台（CUDA 等）即将接入。 +Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架，让大型稀疏 +MoE 模型跑在消费级硬件上——峰值内存由**激活**专家集而非参数量决定。 + +### 核心机制 + +- **SSD 专家 offload**：专家权重按需从存储流式加载，峰值内存由 + 激活集而非参数量决定； +- **prerouter**：训练头提前一步预测专家路由，专家装载与前向计算重叠 + 而非阻塞——解码吞吐**最高 +59%**，收益随存储延迟、模型规模与路由 + 宽度 *K* 增大； +- **Recover-LoRA**：冻结 int4 基模，用 FP teacher 蒸馏训练 LoRA， + 在 4bit 下恢复绝大部分量化损失（见[质量](#质量)）。适配器不合并， + 一份只读基模服务多套适配器。 + +### 平台 + +一个仓库、一套配方、各平台独立 runtime： + +| 平台 | 目录 | 技术栈 | 状态 | +|---|---|---|---| +| **Python**（macOS · Apple Silicon） | [`python/`](python/) | Python + MLX | ✅ 现已可用 | +| **macOS** 桌面 App 与 CLI | [`macos/`](macos/) | Rust | ✅ 已开源（2026-09-30） | +| **iOS** App | [`ios/`](ios/) | Swift + MLX Swift | ✅ 已开源（2026-09-30） | +| **Android** App 与引擎 | [`android/`](android/) | Kotlin + 原生引擎 | ✅ 已开源（2026-09-30） | +| **Windows** | — | — | 🗺️ 路线图中 | + +### 模型 框架随附两个模型档位。每个档位是一个端到端发布：发布的 checkpoint、 训练好的 LoRA 适配器与训练好的 prerouter 头作为整体协同工作。 @@ -35,13 +81,71 @@ Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架。后 适配器文件与 checkpoint 同目录、自动加载，`edge0 serve ` 开箱即跑 训练好的完整管线。 -## 环境要求 +### 设计 + +- **像 transformers 一样使用**：`AutoModel` / `AutoConfig` / `AutoEngine` + 按模型名自动选类； +- **后端隔离（Python 框架）**：在 Python 框架内，全部 MLX 代码收在 + `python/src/edge0/backends/mlx/`，核心逻辑（模型 spec / prerouter / + 流式专家池 / server）只依赖后端门面（`backends/base.py` 的 `core` / + `nn` 门面），新增后端实现同一门面即可平级接入（`backends/cuda/` + 预留插槽），核心代码零改动。iOS / macOS / Android 引擎目前是各自的 + 平台原生技术栈——把所有平台收进统一接入层，正是统一推理框架 + （见[路线图](#-路线图)）要交付的内容； +- **适配器统一为 safetensors**：LoRA 与 prerouter 权重均为带元数据 + （来源、版本、owner 层）的 `.safetensors`，放模型目录或 `artifacts/` + 均可自动解析； +- **模型 + 适配器同目录布局**：一个模型目录同时放基模（`config.json` / + `model*.safetensors` / tokenizer）和该模型的适配器，升级适配器只换 + 适配器文件，基模不动、不 merge。 + +### 质量 + +全部评测由我们使用 [OpenCompass](https://github.com/open-compass/opencompass)、 +在完全相同的设置与参数下对 edge0 模型（int4 + 训练适配器 + prerouter 路由） +与原 fp16 基座模型测得。edge0 管线的损失很小：**edge0-35b 平均仅落后 +3.9 分、edge0-8b 落后 2.8 分**（MMLU-Pro 甚至反超基座）。满分 100： + +| 评测集 | edge0-35b（int4） | Qwen3.6-35B-A3B（fp16） | edge0-8b（int4） | Ling 3.0 tiny（fp16） | +|---|---:|---:|---:|---:| +| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 | +| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 | +| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 | +| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 | +| IFBench | 57.9 | 61.7 | 53.9 | 60.6 | +| **平均** | **79.2** | **83.2** | **69.9** | **72.7** | + +### 性能实测 + +`python/examples/bench.py` 实测（3.3k token prompt prefill → 10 步采样 +warmup → 200 token 计时段，每档 2 轮）： + +| 档位 | 解码速度 | Prefill 吞吐（冷/热）* | 峰值 active 内存 | 测试机器 | +|---|---|---|---|---| +| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB | +| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB | + +*冷 = 进程启动后首请求（专家权重从 SSD 逐页换入）；热 = 后续请求（页缓存常驻）。Prefill 为 ≈3.3k token 长 prompt 的吞吐（`BENCH_LONG=1`）。* + +复现： + +```bash +cd python +python examples/bench.py edge0-35b # 经 $EDGE0_35B_MODEL +python examples/bench.py edge0-8b # 经 $EDGE0_8B_MODEL +``` + +## 🚀 快速开始 + +### Python（macOS · Apple Silicon） + +#### 环境要求 - **系统 / 硬件**：MLX 后端目前仅支持 Apple Silicon 的 macOS - （M1/M2/M3/M4）；CUDA 后端在路线图中，其余平台暂不支持。 + （M1/M2/M3/M4）；CUDA 后端在路线图中，Python 框架暂不支持其他平台。 - **Python**：3.10+（推荐 3.12）。 - **MLX**：`mlx==0.30.6` / `mlx-metal==0.30.6`（`mlx-lm==0.31.0`，见 - `pyproject.toml`）。Apple A18 / A18 Pro 上输出乱码 = mlx 版本旧： + `python/pyproject.toml`）。Apple A18 / A18 Pro 上输出乱码 = mlx 版本旧： `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'` （[#8](https://github.com/Edge0-AI/Edge0/issues/8)）。 - **内存**：短上下文下 `edge0-35b` ≈2.9 GB、`edge0-8b` ≈1.0 GB @@ -50,45 +154,19 @@ Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架。后 - **磁盘**：4bit checkpoint 约 23 GB（`edge0-35b`）/ 4.2 GB （`edge0-8b`）；专家权重 mmap 按需读取，不一次性载入内存。 -## 设计 - -- **像 transformers 一样使用**：`AutoModel` / `AutoConfig` / `AutoEngine` - 按模型名自动选类； -- **后端隔离**：全部 MLX 代码收在 `edge0/backends/mlx/`，核心逻辑 - （模型 spec / prerouter / 流式专家池 / server）只依赖后端门面 - （`edge0/backends/base.py` 的 `core` / `nn` 门面），新增后端实现同一门面 - 即可平级接入（`backends/cuda/` 预留插槽），核心代码零改动； -- **适配器统一为 safetensors**：LoRA 与 prerouter 权重均为带元数据 - （来源、版本、owner 层）的 `.safetensors`，放模型目录或 `artifacts/` - 均可自动解析； -- **模型 + 适配器同目录布局**：一个模型目录同时放基模（`config.json` / - `model*.safetensors` / tokenizer）和该模型的适配器，升级适配器只换适配器文件 - 文件，基模不动、不 merge。 - -## 核心机制 - -- **SSD 专家 offload**：专家权重按需从存储流式加载，峰值内存由 - 激活集而非参数量决定； -- **prerouter**：训练头提前一步预测专家路由，专家装载与前向计算重叠 - 而非阻塞——解码吞吐**最高 +59%**，收益随存储延迟、模型规模与路由 - 宽度 *K* 增大； -- **Recover-LoRA**：冻结 int4 基模，用 FP teacher 蒸馏训练 LoRA， - 在 4bit 下恢复绝大部分量化损失（见[质量](#质量)）。适配器不合并， - 一份只读基模服务多套适配器。 - -## 快速开始 - -### 1) 安装 +#### 1) 安装 ```bash +cd python # Python ≥3.10；MLX 后端需 macOS + Apple Silicon python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]' ``` -### 2) 下载模型 +#### 2) 下载模型 -两个档位发布在 Hugging Face 与 ModelScope——每个仓库把基模 checkpoint 与训练好的 -LoRA + prerouter 适配器打包在**同一目录**，一次下载即为可运行的模型： +两个档位发布在 Hugging Face 与 ModelScope——每个仓库把基模 checkpoint +与训练好的 LoRA + prerouter 适配器打包在**同一目录**，一次下载即为 +可运行的模型： - [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)（约 23 GB） · [ModelScope 镜像](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) - [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)（约 4.2 GB） · [ModelScope 镜像](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview) @@ -114,7 +192,7 @@ models/edge0-35b/ └── prerouter_edge0_35b.safetensors # 训练好的 prerouter 头 ``` -### 3) 指向模型目录 +#### 3) 指向模型目录 档位名经环境变量解析到本地目录（放哪由你决定）： @@ -131,7 +209,7 @@ edge0 demo models/edge0-35b edge0 serve models/edge0-8b ``` -### 4) 运行 +#### 4) 运行 ```bash # 快速演示 @@ -152,7 +230,7 @@ edge0 chat edge0-35b --prompt "用一句话解释流式推理。" `python -m edge0 ...` 等价于 `edge0 ...`。 -### Python API +#### Python API ```python from edge0 import AutoEngine @@ -169,9 +247,9 @@ print(engine._tok.decode(tokens)) engine.close() # 释放 mmap / 专家缓存 ``` -`examples/demo.py` 就是这条最小路径（`edge0 demo` 内部等价运行）。 +`python/examples/demo.py` 就是这条最小路径（`edge0 demo` 内部等价运行）。 -### 模型与适配器 +#### 模型与适配器 - **checkpoint**：原始模型目录（`config.json`、`model*.safetensors`、 tokenizer）。`edge0 serve ` / `AutoEngine.from_pretrained()` @@ -179,68 +257,104 @@ engine.close() # 释放 mmap / 专家缓存 - **适配器**（LoRA + prerouter，safetensors）放两处任一，自动解析： - **模型目录内**（推荐）：与基模同目录，如 `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`； - - `artifacts/`（仓库根，gitignored）：`edge0 convert-adapters` 从 - 训练侧 npz 一次性转换。 + - `artifacts/`（gitignored）：`edge0 convert-adapters` 从训练侧 npz + 一次性转换。 - 发布的模型仓库同时包含基模与当前默认适配器版本， `scripts/fetch_models.py` 下载后即为可运行的模型目录。适配器来源 （训练数据、owner 层分布）见各模型文档页。 - prerouter + LoRA 两条适配器都必需；缺文件时 `edge0` 会给出明确报错 （也可加 `--no-prerouter` / `--no-lora` 直接跑裸基模）。 -## 质量 +#### 文档 -全部评测由我们使用 [OpenCompass](https://github.com/open-compass/opencompass)、 -在完全相同的设置与参数下对 edge0 模型（int4 + 训练适配器 + prerouter 路由） -与原 fp16 基座模型测得。edge0 管线的损失很小：**edge0-35b 平均仅落后 -3.9 分、edge0-8b 落后 2.8 分**（MMLU-Pro 甚至反超基座）。满分 100： +- [架构总览](docs/architecture.md) +- [注意力抽象](docs/attention.md) / [MoE 抽象](docs/moe.md) / [SSD 流式](docs/streaming.md) / [prerouter](docs/prerouter.md) +- [如何接入新模型](docs/adding-a-model.md) +- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) +- 技术报告：[The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063)（[PDF](paper/main.pdf)） -| 评测集 | edge0-35b（int4） | Qwen3.6-35B-A3B（fp16） | edge0-8b（int4） | Ling 3.0 tiny（fp16） | -|---|---:|---:|---:|---:| -| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 | -| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 | -| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 | -| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 | -| IFBench | 57.9 | 61.7 | 53.9 | 60.6 | -| **平均** | **79.2** | **83.2** | **69.9** | **72.7** | +### macOS / iOS / Android -## 性能实测 +三端推理引擎已开源在本仓库，更多细节见各目录 README： -`examples/bench.py` 实测（3.3k token prompt prefill → 10 步采样 warmup → -200 token 计时段，每档 2 轮）： +- **macOS**：本地 CLI / daemon / 桌面 App（Rust）—— 详见 + [`macos/README.md`](macos/README.md) +- **iOS**：iPhone 端侧 App（Swift + MLX Swift）—— 详见 + [`ios/README.md`](ios/README.md) +- **Android**：端侧 App + 原生引擎（Kotlin）—— 详见 + [`android/README.md`](android/README.md) -| 档位 | 解码速度 | Prefill 吞吐（冷/热）* | 峰值 active 内存 | 测试机器 | -|---|---|---|---|---| -| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB | -| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB | +**统一推理框架**（接入层统一，runtime 自动适配 iOS / macOS / Android / +Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#-路线图)。 -*冷 = 进程启动后首请求（专家权重从 SSD 逐页换入）；热 = 后续请求（页缓存常驻）。Prefill 为 ≈3.3k token 长 prompt 的吞吐（`BENCH_LONG=1`）。 +## 🗺️ 路线图 -复现： +- **【2026 年 10 月底】edge0 统一推理框架开源** —— 发布统一推理框架： + **接入层统一**为一套 API（chat / serve / 端侧共用），**runtime 自动 + 适配不同硬件平台**——iOS、macOS、Android、Windows、Python；在本仓库 + 已开源的三端引擎（`ios/` · `macos/` · `android/`）之上统一收编。 +- **Windows** 平台支持（经统一框架）。 +- **CUDA 后端**（Python 框架）—— 插槽已预留在 + `python/src/edge0/backends/cuda/`，核心代码零改动。 +- 更多模型档位与适配器版本在现有管线上持续发布。 -```bash -python examples/bench.py edge0-35b # 经 $EDGE0_35B_MODEL -python examples/bench.py edge0-8b # 经 $EDGE0_8B_MODEL -``` +## 🤝 贡献 -## 验证 +欢迎各种形式的贡献——issue、PR、性能实测报告、新模型接入都算。 + +**Python 框架**（现已可用）： ```bash +cd python +python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]' + pytest # 单元测试（不含真实权重） EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q - # 真实权重生成测试；缺少的档位会明确 skip + # 真实权重生成测试；缺少的档位会明确 skip .venv/bin/python scripts/e2e_smoke.py \ --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b - # staged vs exact 一致性 + 生成冒烟 + # staged vs exact 一致性 + 生成冒烟 scripts/generate_example.py # 完整 API 上手例子（prefill→生成→解码全链路） -examples/demo.py # 最小 API walkthrough（edge0 demo 的等价代码） +examples/demo.py # 最小 API walkthrough（edge0 demo 的等价代码） ``` -## 文档 +每个 PR 都会跑 CI：单元测试（macOS + MLX）与仓库卫生检查（无硬编码 +路径、后端边界、无密钥）。 -- [架构总览](docs/architecture.md) -- [注意力抽象](docs/attention.md) / [MoE 抽象](docs/moe.md) / [SSD 流式](docs/streaming.md) / [prerouter](docs/prerouter.md) -- [如何接入新模型](docs/adding-a-model.md) -- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) +**平台 runtime**（macOS / iOS / Android）：目前是占位目录，贡献指南将 +随统一框架代码一并发布。 + +流程：fork → 功能分支 → 向 `main` 发 PR。请保持卫生检查通过，并为新 +行为补测试。 + +## 📖 引用 + +如果 edge0 对你有帮助，请引用我们的技术报告： + +```bibtex +@article{lin2026other, + title = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD + with Trained Routing Prediction}, + author = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and + Zeng, Xiaodong}, + journal = {arXiv preprint arXiv:2609.18063}, + year = {2026}, + url = {https://arxiv.org/abs/2609.18063} +} +``` + +## 📧 联系我们 + + + +社区与支持渠道即将上线，本节将列出联系我们的官方方式： + +- **邮箱**：_待补充_ +- **Discord**：_待补充_ +- **微信群**：_待补充_ + +Bug 与功能建议请直接提 +[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues)。 ## License From b4d43c11363520c8165db6c24544b01bcaed7826 Mon Sep 17 00:00:00 2001 From: Bin Date: Tue, 29 Sep 2026 17:24:18 +0800 Subject: [PATCH 03/13] fix(cli,models): unify serving port to 8000; repair convert-adapters and demo.py End-to-end testing of the README quick-start surfaced three issues: - edge0 convert-adapters was broken. Its subparser defined no options and cmd_convert re-exec'd the legacy script via runpy without resetting sys.argv, so the subcommand name leaked into the script's argparse ("unrecognized arguments: convert-adapters"), while --npz-dir / --force were rejected by the top-level parser. Define --force / --npz-dir on the subparser and delegate to the legacy main(argv) with a clean argv. - examples/demo.py --model raised FileNotFoundError because tier env-var resolution lived only in the CLI. Resolve $EDGE0__MODEL the same way; an explicit --model-dir still wins, and a missing env var now fails with a clear message. - The edge0-35b / edge0-8b profiles overrode the serving port (8085 / 8083) while the CLI defaulted to 8000. Drop the overrides so every tier inherits the unified port 8000; update the profile assertions and the docs/models/*.md serve/curl examples to match. Also ignore downloaded model checkpoint dirs (/models/, /python/models/) so multi-GB weights are never committed; the pattern is anchored to avoid matching the source package python/src/edge0/models/. Verified: convert-adapters end to end on synthetic npz (2 sources -> 4 safetensors, key normalization, owners, 35b language_model. prefix, idempotent skip, --force rewrite); pytest -m 'not slow' -> 79 passed. Co-Authored-By: Claude Code --- .gitignore | 5 ++++ docs/models/edge0-35b.md | 10 +++---- docs/models/edge0-8b.md | 10 +++---- python/examples/demo.py | 12 ++++++++ python/src/edge0/cli.py | 30 +++++++++++++------ python/src/edge0/models/edge0_35b/__init__.py | 4 +-- python/src/edge0/models/edge0_8b/__init__.py | 4 +-- python/tests/test_registry.py | 4 +-- 8 files changed, 54 insertions(+), 25 deletions(-) diff --git a/.gitignore b/.gitignore index b41702c..623ec53 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,11 @@ dist/ # Trained adapter artifacts (never committed; see README) artifacts/ +# Downloaded model checkpoints (never committed; see README quick start). +# Anchored so the source package python/src/edge0/models/ is not matched. +/models/ +/python/models/ + # macOS .DS_Store diff --git a/docs/models/edge0-35b.md b/docs/models/edge0-35b.md index 99d611b..e6c12af 100644 --- a/docs/models/edge0-35b.md +++ b/docs/models/edge0-35b.md @@ -22,7 +22,7 @@ The performance profile is based on benchmarks of the current release adapter ve | Prefill chunk | 2048 | | Hot window | 4 | | Streaming prefetch history | on (`prefetch_history=True`) | -| Serving port | 8085 | +| Serving port | 8000 | | Measured throughput | 14.9–17.7 tok/s (M4 Pro) | | Measured peak activation memory | ≈ 2.9 GiB | @@ -42,7 +42,7 @@ This tier uses the `LayerOptions.staged_k4()` preset: ### Serving via the CLI ```bash -edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8085 +edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8000 ``` Optional arguments: @@ -99,7 +99,7 @@ cfg = AutoConfig.from_pretrained("/path/to/checkpoint", name="edge0-35b") ### Non-streaming chat ```bash -curl -s http://127.0.0.1:8085/v1/chat/completions \ +curl -s http://127.0.0.1:8000/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{ "model": "edge0-35b", @@ -131,7 +131,7 @@ Optional request fields: `model`, `messages` (with `role`/`content`; content sup ### Streaming chat (requires Flask) ```bash -curl -N http://127.0.0.1:8085/v1/chat/completions \ +curl -N http://127.0.0.1:8000/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{ "model": "edge0-35b", @@ -152,7 +152,7 @@ from edge0 import AutoEngine engine = AutoEngine.from_pretrained( "/path/to/checkpoint", name="edge0-35b", - port=9090, # override the default port 8085 + port=9090, # override the default port 8000 target_tok_s=14.0, # override the acceptance throughput target prerouter=None, # disable the prerouter lora="", # disable LoRA diff --git a/docs/models/edge0-8b.md b/docs/models/edge0-8b.md index 55c2db0..a71943f 100644 --- a/docs/models/edge0-8b.md +++ b/docs/models/edge0-8b.md @@ -22,7 +22,7 @@ The performance profile is based on benchmarks of the current release adapter ve | Prefill chunk | 2048 | | Hot window | 1 | | Streaming prefetch history | on (`prefetch_history=True`) | -| Serving port | 8083 | +| Serving port | 8000 | | Measured throughput | 23.9–25.3 tok/s (M4 Pro) | | Measured peak activation memory | ≈ 1.0 GiB | @@ -41,7 +41,7 @@ This tier uses the `LayerOptions.prod_k8()` preset (aligned with the reference d ### Serving via the CLI ```bash -edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8083 +edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8000 ``` Optional arguments: @@ -99,7 +99,7 @@ cfg = AutoConfig.from_pretrained("/path/to/checkpoint", name="edge0-8b") ### Non-streaming chat ```bash -curl -s http://127.0.0.1:8083/v1/chat/completions \ +curl -s http://127.0.0.1:8000/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{ "model": "edge0-8b", @@ -131,7 +131,7 @@ Optional request fields: `model`, `messages` (with `role`/`content`; content sup ### Streaming chat (requires Flask) ```bash -curl -N http://127.0.0.1:8083/v1/chat/completions \ +curl -N http://127.0.0.1:8000/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{ "model": "edge0-8b", @@ -152,7 +152,7 @@ from edge0 import AutoEngine engine = AutoEngine.from_pretrained( "/path/to/checkpoint", name="edge0-8b", - port=9083, # override the default port 8083 + port=9083, # override the default port 8000 target_tok_s=35.0, # override the acceptance throughput target prerouter=None, # disable the prerouter lora="", # disable LoRA diff --git a/python/examples/demo.py b/python/examples/demo.py index 686700f..236e9b6 100644 --- a/python/examples/demo.py +++ b/python/examples/demo.py @@ -40,6 +40,18 @@ def main() -> int: "already run that path)") args = ap.parse_args() + # ``--model `` resolves through the tier env var (same contract + # as ``edge0 demo``); ``--model-dir`` always wins. + if not args.model_dir and args.model: + import os + + from edge0.cli import TIER_ENV + env = TIER_ENV.get(args.model) + args.model_dir = os.environ.get(env, "") if env else "" + if not args.model_dir: + ap.error(f"--model {args.model} needs ${env or '?'} set to the " + f"checkpoint directory (or pass --model-dir)") + kw = ({"prerouter": None} if args.no_prerouter else demo_kwargs(args.model_dir, args.model)) engine = AutoEngine.from_pretrained(args.model_dir, name=args.model, **kw) diff --git a/python/src/edge0/cli.py b/python/src/edge0/cli.py index d2cbb73..5b3ec54 100644 --- a/python/src/edge0/cli.py +++ b/python/src/edge0/cli.py @@ -209,15 +209,22 @@ def cmd_serve(args) -> int: def cmd_convert(args) -> int: - import runpy - sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve() - .parents[2] / "scripts")) - runpy.run_path( - str(__import__("pathlib").Path(__file__).resolve().parents[2] - / "scripts" / "convert_adapters_legacy.py"), - run_name="__main__", - ) - return 0 + from pathlib import Path + + scripts_dir = Path(__file__).resolve().parents[2] / "scripts" + if str(scripts_dir) not in sys.path: + sys.path.insert(0, str(scripts_dir)) + # Delegate to the legacy converter's own main(argv) with a clean argv, + # so the subcommand name never leaks into its argparse (runpy would + # re-parse sys.argv and choke on "convert-adapters"). + from convert_adapters_legacy import main as _convert_main + + argv: list[str] = [] + if getattr(args, "force", False): + argv.append("--force") + if getattr(args, "npz_dir", None): + argv += ["--npz-dir", args.npz_dir] + return _convert_main(argv) def _add_engine_flags(p) -> None: @@ -290,6 +297,11 @@ def _build_parser() -> argparse.ArgumentParser: p = sub.add_parser("convert-adapters", help="one-shot legacy npz -> safetensors migration") + p.add_argument("--force", action="store_true", + help="reconvert even if the artifact exists") + p.add_argument("--npz-dir", default=None, + help="directory holding the legacy npz exports " + "(default: $EDGE0_NPZ_DIR)") p.set_defaults(fn=cmd_convert) return ap diff --git a/python/src/edge0/models/edge0_35b/__init__.py b/python/src/edge0/models/edge0_35b/__init__.py index 3f27819..0e5bd40 100644 --- a/python/src/edge0/models/edge0_35b/__init__.py +++ b/python/src/edge0/models/edge0_35b/__init__.py @@ -10,7 +10,8 @@ * trained prerouter: 33 heads (owners 6..38), start_layer 7, hidden 512, fp16, consumed directly by the patched MoE block at decode (patch_call=True, cross-token staged decode, K=4). -* serving: port 8085; acceptance ≈13 tok/s, peak active ≈3.3 GB. +* serving: port 8000 (the shared CLI default); acceptance ≈13 tok/s, + peak active ≈3.3 GB. """ from __future__ import annotations @@ -64,7 +65,6 @@ def _defaults(cls, model_dir: str) -> "Qwen35Config": hot_window=4, intra_staging=False, prefetch_history=True, - port=8085, target_tok_s=13.0, peak_active_mem_mb=3400.0, ) diff --git a/python/src/edge0/models/edge0_8b/__init__.py b/python/src/edge0/models/edge0_8b/__init__.py index 7fcc724..5453580 100644 --- a/python/src/edge0/models/edge0_8b/__init__.py +++ b/python/src/edge0/models/edge0_8b/__init__.py @@ -11,7 +11,8 @@ (explicit owners 1..22), start_layer 1, hidden 512, fp16, feature_topk "executed", consumed inside ``BailingSparseMoE`` from ``prerouter_cache`` logits (patch_call=False). -* serving: port 8083; acceptance ≈33 tok/s, peak active ≈1.4 GB. +* serving: port 8000 (the shared CLI default); acceptance ≈33 tok/s, + peak active ≈1.4 GB. """ from __future__ import annotations @@ -72,7 +73,6 @@ def _defaults(cls, model_dir: str) -> "Ling8BConfig": hot_window=1, intra_staging=False, prefetch_history=True, - port=8083, target_tok_s=33.0, peak_active_mem_mb=1400.0, ) diff --git a/python/tests/test_registry.py b/python/tests/test_registry.py index 9cf7932..d3a61d9 100644 --- a/python/tests/test_registry.py +++ b/python/tests/test_registry.py @@ -65,7 +65,7 @@ def test_qwen35_profile(): name="edge0-35b", history_slots=True).history_slots is True assert cfg.gen.temperature == 0.7 assert 248046 in cfg.gen.eos_ids - assert cfg.port == 8085 + assert cfg.port == 8000 # unified serving port (CLI default) def test_ling8b_profile(): @@ -82,7 +82,7 @@ def test_ling8b_profile(): assert cfg.prerouter.patch_call is False assert cfg.gen.repetition_penalty == 1.1 assert 156895 in cfg.gen.eos_ids - assert cfg.port == 8083 + assert cfg.port == 8000 # unified serving port (CLI default) def test_override_and_reject(): From e764992584107eb2ce24026f02dba6a6dd6c1663 Mon Sep 17 00:00:00 2001 From: Bin Date: Tue, 29 Sep 2026 17:25:33 +0800 Subject: [PATCH 04/13] docs(readme): polish EN/ZH READMEs and fix roadmap anchors - Drop the section-heading emoji (News / About / Getting Started / Roadmap / Contributing / Citation / Contact Us) and repair the in-repo anchors that the removal invalidated (#-roadmap -> #roadmap) in both root READMEs and the macos / ios / android placeholder READMEs. - Add the required "model" field to the /v1/chat/completions curl example; the server rejects requests without it (surfaced during testing). - Fill in the Contact Us email; sync the Chinese README with all of the above. Co-Authored-By: Claude Code --- README.md | 26 ++++++++++++-------------- README_zh.md | 28 +++++++++++++--------------- android/README.md | 2 +- ios/README.md | 2 +- macos/README.md | 2 +- 5 files changed, 28 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 110e96c..0f6ba4f 100644 --- a/README.md +++ b/README.md @@ -20,14 +20,14 @@ English | [中文](README_zh.md) -## 📰 News +## News -- **[2026-09-30]** We released the **edge0 inference engines for three platforms — iOS, macOS and Android** — so users get the best inference experience across architectures and platforms. The source is open-sourced in this repo ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/)) — see each directory's README for details. The **unified inference framework** follows by the **end of October 2026**; see the [Roadmap](#-roadmap). +- **[2026-09-30]** We released the **edge0 inference engines for three platforms — iOS, macOS and Android** — so users get the best inference experience across architectures and platforms. The source is open-sourced in this repo ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/)) — see each directory's README for details. The **unified inference framework** follows by the **end of October 2026**; see the [Roadmap](#roadmap). - **[2026-09-20]** `--prefill-ondemand` lands for `edge0 demo / chat / serve` ([#112](https://github.com/Edge0-AI/Edge0/pull/112)). - **[2026-09-16]** Our technical report is on arXiv: [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063). - **[2026-09-08]** Initial open-source release of **edge0**, together with both model tiers — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) and [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — on Hugging Face and ModelScope. -## 🤔 About +## About **edge0** is an open-source streaming MoE inference framework. It generalizes the production-proven recipe — **SSD expert offload + @@ -90,7 +90,7 @@ are co-located with each checkpoint and load automatically, so zero changes to core code. The iOS / macOS / Android engines ship platform-native stacks today — bringing every platform under one access layer is exactly what the unified inference framework - (see [Roadmap](#-roadmap)) will deliver; + (see [Roadmap](#roadmap)) will deliver; - **Adapters as safetensors**: LoRA and prerouter weights are `.safetensors` files with provenance metadata (source, version, owner layers), resolved from the model directory or `artifacts/`; @@ -138,7 +138,7 @@ python examples/bench.py edge0-35b # via $EDGE0_35B_MODEL python examples/bench.py edge0-8b # via $EDGE0_8B_MODEL ``` -## 🚀 Getting Started +## Getting Started ### Python (macOS · Apple Silicon) @@ -226,7 +226,7 @@ edge0 serve edge0-35b ```bash curl http://127.0.0.1:8000/v1/chat/completions \ -H 'Content-Type: application/json' \ - -d '{"messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}' + -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}' # 5) One-shot chat (pass --max-new to cap length; add --show-thinking to # print the model's reasoning block too) @@ -294,9 +294,9 @@ in each directory's README: The **unified inference framework** — one access layer, runtime auto-adapting to iOS / macOS / Android / Windows / Python — arrives by -the **end of October 2026**; see the [Roadmap](#-roadmap). +the **end of October 2026**; see the [Roadmap](#roadmap). -## 🗺️ Roadmap +## Roadmap - **[End of Oct 2026] edge0 unified inference framework** — we will open-source a unified inference framework: **one unified access @@ -310,7 +310,7 @@ the **end of October 2026**; see the [Roadmap](#-roadmap). `python/src/edge0/backends/cuda/`, core code needs zero changes. - More model tiers and adapter releases on the existing pipeline. -## 🤝 Contributing +## Contributing Contributions are welcome — issues, PRs, benchmark reports and model ports all count. @@ -341,7 +341,7 @@ with the unified framework code. Workflow: fork → feature branch → PR against `main`. Please keep the hygiene suite green and add tests for new behavior. -## 📖 Citation +## Citation If you find edge0 useful, please cite our technical report: @@ -357,16 +357,14 @@ If you find edge0 useful, please cite our technical report: } ``` -## 📧 Contact Us +## Contact Us Community and support channels are coming soon — this section will list the official ways to reach us: -- **Email**: _to be added_ -- **Discord**: _to be added_ -- **WeChat group**: _to be added_ +- **Email**: samuel@edge0.ai For bugs and feature requests, please use [GitHub Issues](https://github.com/Edge0-AI/Edge0/issues). diff --git a/README_zh.md b/README_zh.md index 6631ef3..aeb9ce6 100644 --- a/README_zh.md +++ b/README_zh.md @@ -20,13 +20,13 @@ -## 📰 新闻 +## 新闻 - **【2026-09-30】** 我们发布了**三端推理引擎**，适配 iOS、macOS、Android 平台，让用户在不同架构、不同平台上都有最佳的推理体验。三端源码已 开源至本仓库（[`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/)），更多细节见各目录 README；**统一推理框架** - 将于 **2026 年 10 月底**发布，详见[路线图](#-路线图)。 + 将于 **2026 年 10 月底**发布，详见[路线图](#路线图)。 - **【2026-09-20】** `edge0 demo / chat / serve` 新增 `--prefill-ondemand` （[#112](https://github.com/Edge0-AI/Edge0/pull/112)）。 - **【2026-09-16】** 技术报告上线 arXiv：[The Other Half of the Memory @@ -37,7 +37,7 @@ 与 [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) 同步登陆 Hugging Face 与 ModelScope。 -## 🤔 关于 edge0 +## 关于 edge0 **edge0** 是一个开源的流式 MoE 推理框架：把「SSD 专家 offload + Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架，让大型稀疏 @@ -91,7 +91,7 @@ MoE 模型跑在消费级硬件上——峰值内存由**激活**专家集而非 `nn` 门面），新增后端实现同一门面即可平级接入（`backends/cuda/` 预留插槽），核心代码零改动。iOS / macOS / Android 引擎目前是各自的 平台原生技术栈——把所有平台收进统一接入层，正是统一推理框架 - （见[路线图](#-路线图)）要交付的内容； + （见[路线图](#路线图)）要交付的内容； - **适配器统一为 safetensors**：LoRA 与 prerouter 权重均为带元数据 （来源、版本、owner 层）的 `.safetensors`，放模型目录或 `artifacts/` 均可自动解析； @@ -135,7 +135,7 @@ python examples/bench.py edge0-35b # 经 $EDGE0_35B_MODEL python examples/bench.py edge0-8b # 经 $EDGE0_8B_MODEL ``` -## 🚀 快速开始 +## 快速开始 ### Python（macOS · Apple Silicon） @@ -222,7 +222,7 @@ edge0 serve edge0-35b ```bash curl http://127.0.0.1:8000/v1/chat/completions \ -H 'Content-Type: application/json' \ - -d '{"messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}' + -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}' # 5) 单轮对话（--max-new 限制生成长度；加 --show-thinking 会一并打印思考块） edge0 chat edge0-35b --prompt "用一句话解释流式推理。" @@ -285,9 +285,9 @@ engine.close() # 释放 mmap / 专家缓存 [`android/README.md`](android/README.md) **统一推理框架**（接入层统一，runtime 自动适配 iOS / macOS / Android / -Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#-路线图)。 +Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#路线图)。 -## 🗺️ 路线图 +## 路线图 - **【2026 年 10 月底】edge0 统一推理框架开源** —— 发布统一推理框架： **接入层统一**为一套 API（chat / serve / 端侧共用），**runtime 自动 @@ -298,7 +298,7 @@ Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#-路线 `python/src/edge0/backends/cuda/`，核心代码零改动。 - 更多模型档位与适配器版本在现有管线上持续发布。 -## 🤝 贡献 +## 贡献 欢迎各种形式的贡献——issue、PR、性能实测报告、新模型接入都算。 @@ -327,7 +327,7 @@ examples/demo.py # 最小 API walkthrough（edge0 demo 的等价代 流程：fork → 功能分支 → 向 `main` 发 PR。请保持卫生检查通过，并为新 行为补测试。 -## 📖 引用 +## 引用 如果 edge0 对你有帮助，请引用我们的技术报告： @@ -343,15 +343,13 @@ examples/demo.py # 最小 API walkthrough（edge0 demo 的等价代 } ``` -## 📧 联系我们 +## 联系我们 - + 社区与支持渠道即将上线，本节将列出联系我们的官方方式： -- **邮箱**：_待补充_ -- **Discord**：_待补充_ -- **微信群**：_待补充_ +- **邮箱**：samuel@edge0.ai Bug 与功能建议请直接提 [GitHub Issues](https://github.com/Edge0-AI/Edge0/issues)。 diff --git a/android/README.md b/android/README.md index 6068ea4..9be16bd 100644 --- a/android/README.md +++ b/android/README.md @@ -11,7 +11,7 @@ Recover-LoRA, and prerouter routing prediction. The next milestone is the **unified inference framework** — one access layer, runtime auto-adapting to the hardware platform (iOS / macOS / Android / Windows / Python) — targeted for the **end of October 2026**; -see the [roadmap](../README.md#-roadmap). +see the [roadmap](../README.md#roadmap). See the [repository README](../README.md) for the multi-platform picture. diff --git a/ios/README.md b/ios/README.md index 389fa70..7fc05a3 100644 --- a/ios/README.md +++ b/ios/README.md @@ -11,7 +11,7 @@ Recover-LoRA, and prerouter routing prediction. The next milestone is the **unified inference framework** — one access layer, runtime auto-adapting to the hardware platform (iOS / macOS / Android / Windows / Python) — targeted for the **end of October 2026**; -see the [roadmap](../README.md#-roadmap). +see the [roadmap](../README.md#roadmap). See the [repository README](../README.md) for the multi-platform picture. diff --git a/macos/README.md b/macos/README.md index 35e7ddd..e4efb5c 100644 --- a/macos/README.md +++ b/macos/README.md @@ -12,7 +12,7 @@ and serves them locally behind an OpenAI-compatible API. The next milestone is the **unified inference framework** — one access layer, runtime auto-adapting to the hardware platform (iOS / macOS / Android / Windows / Python) — targeted for the **end of October 2026**; -see the [roadmap](../README.md#-roadmap). +see the [roadmap](../README.md#roadmap). See the [repository README](../README.md) for the multi-platform picture. From 27866bd11974cf93f62d639eb67af8d85d378bae Mon Sep 17 00:00:00 2001 From: Bin Date: Tue, 29 Sep 2026 17:47:50 +0800 Subject: [PATCH 05/13] docs(readme): fold the Q4 2026 model & algorithm roadmap into Roadmap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split the Roadmap into "Platforms & systems" (the existing unified framework / Windows / CUDA items) and "Models & algorithms — Q4 2026", summarizing the algorithm team's Q4 plan for a public audience: - Next-gen architecture support (Qwen3.8-Flash class): hybrid linear attention (GDN + QSA), gated multi-branch residual, and N-gram embedding — designs that suit SSD streaming offload (O(1)-state attention, lookup-only N-gram tables). - Latent thinking + batched expert pre-prediction: move expert routing from per-position to once per block so load volume decouples from the reasoning loop count, with cross-block prefetch; measured as end-to-end thinking-phase time at matched accuracy. Internal research detail (training configs, supervision ablations, quantified speedup targets, weight-availability contingencies) is intentionally omitted. EN and ZH READMEs updated in sync. Co-Authored-By: Claude Code --- README.md | 24 ++++++++++++++++++++++++ README_zh.md | 19 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/README.md b/README.md index 0f6ba4f..f1cb73c 100644 --- a/README.md +++ b/README.md @@ -298,6 +298,8 @@ the **end of October 2026**; see the [Roadmap](#roadmap). ## Roadmap +**Platforms & systems** + - **[End of Oct 2026] edge0 unified inference framework** — we will open-source a unified inference framework: **one unified access layer** (a single API across chat / serve / on-device use), with the @@ -308,6 +310,28 @@ the **end of October 2026**; see the [Roadmap](#roadmap). - **Windows** platform support (via the unified framework). - **CUDA backend** for the Python framework — reserved slot at `python/src/edge0/backends/cuda/`, core code needs zero changes. + +**Models & algorithms — Q4 2026** + +Q4 works two fronts: bringing a next-generation architecture into the +framework, and turning latent reasoning into a real latency saving rather +than just an arithmetic one. + +- **Next-gen architecture support (Qwen3.8-Flash class)** — run hybrid + linear attention (GDN + QSA), gated multi-branch residual, and N-gram + embedding topologies on edge0. These designs suit SSD streaming offload + naturally: O(1)-state attention keeps long thinking from becoming a + KV-cache problem, and lookup-only N-gram tables stream on demand. Goal: + the tier runs on a single device and benchmarks within an acceptable + gap of the fp16 base. +- **Latent thinking + batched expert pre-prediction** — make latent + reasoning a *latency* saving, not only an arithmetic one. The core + engineering problem: move expert routing from per-position to **once + per block**, so one prediction covers every position and round of a + block and **expert load volume decouples from the reasoning loop + count** — plus cross-block prefetch that loads the next block's experts + inside the current block's compute window. Progress is measured as + **end-to-end thinking-phase time at matched accuracy** (never tokens/s). - More model tiers and adapter releases on the existing pipeline. ## Contributing diff --git a/README_zh.md b/README_zh.md index aeb9ce6..37f615f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -289,6 +289,8 @@ Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#路线 ## 路线图 +**平台与系统** + - **【2026 年 10 月底】edge0 统一推理框架开源** —— 发布统一推理框架： **接入层统一**为一套 API（chat / serve / 端侧共用），**runtime 自动 适配不同硬件平台**——iOS、macOS、Android、Windows、Python；在本仓库 @@ -296,6 +298,23 @@ Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#路线 - **Windows** 平台支持（经统一框架）。 - **CUDA 后端**（Python 框架）—— 插槽已预留在 `python/src/edge0/backends/cuda/`，核心代码零改动。 + +**模型与算法 —— 2026 Q4** + +Q4 推进两条主线：把下一代架构引入框架，以及把潜在推理（latent +thinking）变成真正的延迟收益、而不只是算力节省。 + +- **下一代架构支持（Qwen3.8-Flash 级）** —— 在 edge0 上跑通混合线性 + 注意力（GDN + QSA）、门控多分支残差、N-gram embedding 拓扑。这些 + 设计天然契合 SSD 流式 offload：O(1) 状态的注意力层让长思考不再是 + KV cache 问题，纯查表的 N-gram 表按需流式加载。目标：该档位在单 + 设备跑通，且基准与 fp16 基座的差距在可接受范围内。 +- **潜在思考 + 批量专家预预测** —— 让潜在推理成为*延迟*收益，而不 + 只是算力收益。核心工程问题：把专家路由从 per-position 改为**每 block + 一次**，一次预测覆盖一个 block 内所有位置与所有轮次，使**专家加载量 + 与推理循环次数解耦**；配合跨 block 预取，在当前 block 的计算窗口内 + 加载下一个 block 的专家集。进展以**等准确率下的端到端思考阶段耗时** + 衡量（而非 tokens/s）。 - 更多模型档位与适配器版本在现有管线上持续发布。 ## 贡献 From d464ceb737f12bac6c14e59b5e03136250347db2 Mon Sep 17 00:00:00 2001 From: Bin Date: Tue, 29 Sep 2026 18:09:29 +0800 Subject: [PATCH 06/13] clean up convert_adapters command --- README.md | 4 +- README_zh.md | 4 +- docs/adding-a-model.md | 4 +- docs/architecture.md | 8 +- python/scripts/convert_adapters_legacy.py | 159 ---------------------- python/src/edge0/adapters/lora.py | 4 +- python/src/edge0/cli.py | 29 ---- python/src/edge0/models/base.py | 8 +- 8 files changed, 14 insertions(+), 206 deletions(-) delete mode 100644 python/scripts/convert_adapters_legacy.py diff --git a/README.md b/README.md index f1cb73c..1d5b87b 100644 --- a/README.md +++ b/README.md @@ -265,8 +265,8 @@ runs this exact path). location automatically: - the model directory (recommended): side by side with the base, e.g. `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`; - - `artifacts/` (gitignored): convert once from training-side npz - exports via `edge0 convert-adapters --npz-dir ...`. + - `artifacts/` at the Python project root (gitignored): an optional + fallback cache for adapter safetensors not co-located with the model. - The published model repos bundle both the base checkpoint and the current default adapter release, so `scripts/fetch_models.py` produces a ready-to-run model directory. Check each model's doc page for its diff --git a/README_zh.md b/README_zh.md index 37f615f..6b76609 100644 --- a/README_zh.md +++ b/README_zh.md @@ -257,8 +257,8 @@ engine.close() # 释放 mmap / 专家缓存 - **适配器**（LoRA + prerouter，safetensors）放两处任一，自动解析： - **模型目录内**（推荐）：与基模同目录，如 `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`； - - `artifacts/`（gitignored）：`edge0 convert-adapters` 从训练侧 npz - 一次性转换。 + - `artifacts/`（Python 项目根，gitignored）：可选的回退缓存，放置未 + 与模型同目录的适配器 safetensors。 - 发布的模型仓库同时包含基模与当前默认适配器版本， `scripts/fetch_models.py` 下载后即为可运行的模型目录。适配器来源 （训练数据、owner 层分布）见各模型文档页。 diff --git a/docs/adding-a-model.md b/docs/adding-a-model.md index f14ac8c..be304bb 100644 --- a/docs/adding-a-model.md +++ b/docs/adding-a-model.md @@ -42,7 +42,7 @@ The subclass only needs to implement the class method `_defaults(model_dir) -> C Key points: -- The Python project root (`python/`) ships `scripts/convert_adapters_legacy.py`, which one-shot converts the training npz exports into safetensors artifacts under `artifacts/`; `ModelConfig.artifact(name)` (`models/base.py:28–30`) returns the absolute paths of those artifacts, and the adapter uses them to fill in the LoRA / prerouter weight paths. +- `ModelConfig.artifact(name)` (`models/base.py`) resolves an adapter's absolute path — from the model directory first, falling back to the gitignored `artifacts/` cache — and the adapter uses it to fill in the LoRA / prerouter weight paths. - LoRA overrides go through `resolve_lora` (`models/base.py:88–95`): a bare model name resolves to that tier's artifacts, `"model_dir"` means "keep the training weights in place", and an empty string disables it. ## Step 3: `build_model` / `build_engine` @@ -105,7 +105,7 @@ The lookup order of `_resolve_name` (registry.py:42–56): explicit `name` → t - `config.json` exists, and its `model_type` is registered (or covered by your `TYPE_ALIASES`); - the weights are in safetensors format, with key prefixes matching `moe_spec.key_template` (e.g. `language_model.model.layers.N.mlp.switch_mlp`); -- the LoRA / prerouter weights are `.safetensors` files carrying metadata (see the comment at the top of `models/base.py`: they are converted one-shot from the training npz by `convert_adapters_legacy.py`). +- the LoRA / prerouter weights are `.safetensors` files carrying metadata (see the comment at the top of `models/base.py`). The directory basename is the last-resort resolution fallback (`_model_type_from_dir`), so aligning the directory name with the registered name is recommended, but not required. diff --git a/docs/architecture.md b/docs/architecture.md index 3a07c42..c5adc22 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -19,12 +19,10 @@ box: `edge0-35b` (Qwen3.6-35B-A3B, K=4) and `edge0-8b` (Ling 3.0 hybrid, K=8). future CUDA backend (`backends/cuda/`, selected via the `EDGE0_BACKEND` environment variable); 3. **Adapters unified on safetensors**: LoRA and prerouter weights are both - `.safetensors` files carrying metadata; legacy npz training exports are - converted by a one-off migration script and then deprecated; + `.safetensors` files carrying metadata; 4. **Unified terminology**: the pre-routing head is always called the - **prerouter** — zero legacy-term residue in code and docs. The legacy npz - exports are migrated by `scripts/convert_adapters_legacy.py`, which - normalizes their key namespace to `layers...weight`. + **prerouter** — zero legacy-term residue in code and docs. Prerouter + weight keys are normalized to `layers...weight`. ## Layered Structure diff --git a/python/scripts/convert_adapters_legacy.py b/python/scripts/convert_adapters_legacy.py deleted file mode 100644 index 09c1b96..0000000 --- a/python/scripts/convert_adapters_legacy.py +++ /dev/null @@ -1,159 +0,0 @@ -#!/usr/bin/env python3 -"""One-shot legacy adapter converter: npz -> safetensors. - -The training pipelines exported adapters as ``.npz`` (qwen round-7 and -the training side). edge0 consumes safetensors ONLY — this script converts each -legacy pair once into ``artifacts/`` (gitignored), writing the -conversion provenance (model, kind, K, r/alpha, source md5, format -version) into the safetensors ``__metadata__``. The npz sources are -NOT deleted (see the conversion note in README); they are only needed -to regenerate the artifacts. - -Usage: - python scripts/convert_adapters_legacy.py [--force] - -Idempotent: existing artifacts are skipped unless ``--force``. -""" - -from __future__ import annotations - -import argparse -import hashlib -import os -import json -import sys -from datetime import datetime, timezone -from pathlib import Path - -import numpy as np - -ROOT = Path(__file__).resolve().parents[1] -ARTIFACTS = ROOT / "artifacts" - -# (npz path, tier, kind) — every legacy source this script understands. -# Training-side npz exports discovered per machine; pass --npz-dir (or -# edit this list) to point at your own exports. -NPZ_DIRS = [ - Path(os.environ.get("EDGE0_NPZ_DIR", "/path/to/npz/exports")), -] -SOURCES = [ - ("prerouter_.npz", "edge0-35b", "prerouter"), - ("lora_.npz", "edge0-35b", "lora"), - ("prerouter_.npz", "edge0-8b", "prerouter"), - ("lora_.npz", "edge0-8b", "lora"), -] - - -def _find_source(name: str) -> Path | None: - for d in NPZ_DIRS: - p = d / name - if p.is_file(): - return p - return None - - -def md5(path: Path) -> str: - h = hashlib.md5() - with open(path, "rb") as f: - for chunk in iter(lambda: f.read(1 << 20), b""): - h.update(chunk) - return h.hexdigest() - - -def rewrite_key(key: str, tier: str, kind: str) -> str: - """Map a legacy npz key to the edge0 artifact key. - - prerouter heads are normalized to ``layers...weight``; - qwen lora gains the ``language_model.`` prefix (the checkpoint - namespace the adapter is applied into). - """ - if kind == "prerouter": - # Some families use layers.N.*; others namespace the layer index - # with a prefix segment (...weight). Both are - # normalized to layers...weight. - parts = key.split(".") - if (len(parts) >= 4 and parts[0] != "layers" - and parts[1].isdigit()): - return f"layers.{parts[1]}.{parts[2]}.{parts[3]}" - return key - if kind == "lora" and tier == "edge0-35b" and key.startswith("model.layers."): - return "language_model." + key - return key - - -def convert(source: Path, tier: str, kind: str, force: bool = False) -> Path: - name = { - ("edge0-35b", "prerouter"): "prerouter_edge0_35b_k4.safetensors", - ("edge0-35b", "lora"): "lora_edge0_35b_k4.safetensors", - ("edge0-8b", "prerouter"): "prerouter_edge0_8b.safetensors", - ("edge0-8b", "lora"): "lora_edge0_8b.safetensors", - }[(tier, kind)] - out = ARTIFACTS / name - if out.exists() and not force: - print(f"skip {out.name} (exists)") - return out - - data = np.load(source) - tensors = {} - for key in data.files: - tensors[rewrite_key(key, tier, kind)] = data[key] - - # provenance: rank from the A matrix, alpha from the training run - # (32 for both shipped pairs), K from the tier profile. - meta = { - "model": tier, - "kind": kind, - "K": "4" if tier == "edge0-35b" else "8", - "r": "16", "alpha": "32", - "source": str(source), - "source_md5": md5(source), - "converted": datetime.now(timezone.utc).isoformat(timespec="seconds"), - "format_version": "1", - } - if kind == "prerouter": - owners = sorted({int(k.split(".")[1]) - for k in tensors if k.startswith("layers.")}) - meta["owners"] = json.dumps(owners) - - from safetensors.numpy import save_file as save_np - save_np(tensors, str(out), metadata={"__metadata__": json.dumps(meta)}) - print(f"write {out.name} ({len(tensors)} tensors, " - f"{out.stat().st_size >> 20} MiB)") - return out - - -def main(argv: list[str] | None = None) -> int: - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--force", action="store_true", - help="reconvert even if the artifact exists") - ap.add_argument("--npz-dir", default=None, - help="directory holding the legacy npz exports " - "(default: $EDGE0_NPZ_DIR)") - args = ap.parse_args(argv) - - if args.npz_dir: - NPZ_DIRS.insert(0, Path(args.npz_dir)) - ARTIFACTS.mkdir(parents=True, exist_ok=True) - missing = [] - resolved = [] - for name, tier, kind in SOURCES: - p = _find_source(name) - if p is None: - missing.append(name) - else: - resolved.append((p, tier, kind)) - if missing: - print("missing legacy sources (skip their conversion):") - for m in missing: - print(" -", m) - for path, tier, kind in resolved: - try: - convert(path, tier, kind, force=args.force) - except Exception as exc: # noqa: BLE001 — report and continue - print(f"FAIL {path.name}: {exc}", file=sys.stderr) - return 1 - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/python/src/edge0/adapters/lora.py b/python/src/edge0/adapters/lora.py index b13bf32..1103dd8 100644 --- a/python/src/edge0/adapters/lora.py +++ b/python/src/edge0/adapters/lora.py @@ -16,9 +16,7 @@ where ```` is the module path inside the loaded model (e.g. ``model.layers.3.attention.q_proj`` or -``language_model.model.layers.0.linear_attn.in_proj_qkv``). The legacy -npz files used the same A/B naming, so conversion is a pure format -migration (see ``scripts/convert_adapters_legacy.py``). +``language_model.model.layers.0.linear_attn.in_proj_qkv``). """ from __future__ import annotations diff --git a/python/src/edge0/cli.py b/python/src/edge0/cli.py index 5b3ec54..097442c 100644 --- a/python/src/edge0/cli.py +++ b/python/src/edge0/cli.py @@ -6,7 +6,6 @@ edge0 serve start the HTTP server (one model, queued generations) edge0 chat one-shot prompt -> answer on the terminal edge0 models list registered tiers and their default profiles - edge0 convert-adapters one-shot legacy npz -> safetensors migration """ from __future__ import annotations @@ -208,25 +207,6 @@ def cmd_serve(args) -> int: return 0 -def cmd_convert(args) -> int: - from pathlib import Path - - scripts_dir = Path(__file__).resolve().parents[2] / "scripts" - if str(scripts_dir) not in sys.path: - sys.path.insert(0, str(scripts_dir)) - # Delegate to the legacy converter's own main(argv) with a clean argv, - # so the subcommand name never leaks into its argparse (runpy would - # re-parse sys.argv and choke on "convert-adapters"). - from convert_adapters_legacy import main as _convert_main - - argv: list[str] = [] - if getattr(args, "force", False): - argv.append("--force") - if getattr(args, "npz_dir", None): - argv += ["--npz-dir", args.npz_dir] - return _convert_main(argv) - - def _add_engine_flags(p) -> None: """Engine-construction flags shared by demo / chat / serve.""" p.add_argument("--no-prerouter", action="store_true") @@ -295,15 +275,6 @@ def _build_parser() -> argparse.ArgumentParser: _add_engine_flags(p) p.set_defaults(fn=cmd_serve) - p = sub.add_parser("convert-adapters", - help="one-shot legacy npz -> safetensors migration") - p.add_argument("--force", action="store_true", - help="reconvert even if the artifact exists") - p.add_argument("--npz-dir", default=None, - help="directory holding the legacy npz exports " - "(default: $EDGE0_NPZ_DIR)") - p.set_defaults(fn=cmd_convert) - return ap diff --git a/python/src/edge0/models/base.py b/python/src/edge0/models/base.py index bb1d73d..d6ede85 100644 --- a/python/src/edge0/models/base.py +++ b/python/src/edge0/models/base.py @@ -19,16 +19,16 @@ from edge0.prerouter.spec import PrerouterSpec from edge0.streaming.options import LayerOptions -# Project-root artifacts (gitignored; produced once by -# scripts/convert_adapters_legacy.py from the training npz exports, then -# the source npz are discarded). +# Project-root artifacts (gitignored): an optional local cache of adapter +# safetensors. Adapters normally ship inside the model directory; this is +# only the fallback location artifact() returns when they are not there. ARTIFACTS_DIR = Path(__file__).resolve().parents[3] / "artifacts" def artifact(name: str, model_dir: str | None = None) -> str: """Path of one adapter file: the model's own directory first (model and its adapters side by side in one - directory), falling back to ``artifacts/`` (conversion cache).""" + directory), falling back to ``artifacts/`` (optional local cache).""" if model_dir: cand = Path(model_dir) / name if cand.is_file(): From aa81008089535300f3707e2340d54a99acf41103 Mon Sep 17 00:00:00 2001 From: Bin Date: Wed, 30 Sep 2026 11:23:08 +0800 Subject: [PATCH 07/13] adjust readme --- README.md | 26 +++++++++++++------------- README_zh.md | 35 ++++++++++++++++++----------------- android/README.md | 4 ++-- ios/README.md | 4 ++-- macos/README.md | 6 +++--- windows/README.md | 25 +++++++++++++++++++++++++ 6 files changed, 63 insertions(+), 37 deletions(-) create mode 100644 windows/README.md diff --git a/README.md b/README.md index 1d5b87b..845e8d2 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ English | [中文](README_zh.md) ## News -- **[2026-09-30]** We released the **edge0 inference engines for three platforms — iOS, macOS and Android** — so users get the best inference experience across architectures and platforms. The source is open-sourced in this repo ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/)) — see each directory's README for details. The **unified inference framework** follows by the **end of October 2026**; see the [Roadmap](#roadmap). +- **[2026-09-30]** We released the **edge0 inference engines for four platforms — iOS, macOS, Android and Windows** — so users get the best inference experience across architectures and platforms. The source is open-sourced in this repo ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/) · [`windows/`](windows/)) — see each directory's README for details. The **unified inference framework** follows in **Q4 2026**; see the [Roadmap](#roadmap). - **[2026-09-20]** `--prefill-ondemand` lands for `edge0 demo / chat / serve` ([#112](https://github.com/Edge0-AI/Edge0/pull/112)). - **[2026-09-16]** Our technical report is on arXiv: [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063). - **[2026-09-08]** Initial open-source release of **edge0**, together with both model tiers — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) and [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — on Hugging Face and ModelScope. @@ -59,7 +59,7 @@ One repo, one recipe, per-platform runtimes: | **macOS** app & CLI | [`macos/`](macos/) | Rust | ✅ Open-sourced (2026-09-30) | | **iOS** app | [`ios/`](ios/) | Swift + MLX Swift | ✅ Open-sourced (2026-09-30) | | **Android** app & engine | [`android/`](android/) | Kotlin + native engine | ✅ Open-sourced (2026-09-30) | -| **Windows** | — | — | 🗺️ On the roadmap | +| **Windows** app & engine | [`windows/`](windows/) | C++ + Vulkan | ✅ Open-sourced (2026-09-30) | ### Models @@ -283,35 +283,37 @@ runs this exact path). - [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) - Technical report: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](paper/main.pdf)) -### macOS / iOS / Android +### macOS / iOS / Android / Windows -The three platform engines are open-sourced in this repo — more details +The four platform engines are open-sourced in this repo — more details in each directory's README: - **macOS**: local CLI / daemon / desktop app (Rust) — see [`macos/README.md`](macos/README.md) - **iOS**: on-device iPhone app (Swift + MLX Swift) — see [`ios/README.md`](ios/README.md) - **Android**: on-device app + native engine (Kotlin) — see [`android/README.md`](android/README.md) +- **Windows**: desktop app + native engine (C++ + Vulkan) — see [`windows/README.md`](windows/README.md) The **unified inference framework** — one access layer, runtime -auto-adapting to iOS / macOS / Android / Windows / Python — arrives by -the **end of October 2026**; see the [Roadmap](#roadmap). +auto-adapting to iOS / macOS / Android / Windows / Python — arrives in +**Q4 2026**; see the [Roadmap](#roadmap). ## Roadmap +### Q4 2026 + **Platforms & systems** -- **[End of Oct 2026] edge0 unified inference framework** — we will +- **edge0 unified inference framework** — we will open-source a unified inference framework: **one unified access layer** (a single API across chat / serve / on-device use), with the **runtime automatically adapting to the hardware platform** — iOS, - macOS, Android, Windows and Python. It builds on the three platform + macOS, Android, Windows and Python. It builds on the platform engines already open-sourced in this repo (`ios/` · `macos/` · - `android/`). -- **Windows** platform support (via the unified framework). + `android/` · `windows/`). - **CUDA backend** for the Python framework — reserved slot at `python/src/edge0/backends/cuda/`, core code needs zero changes. -**Models & algorithms — Q4 2026** +**Models & algorithms** Q4 works two fronts: bringing a next-generation architecture into the framework, and turning latent reasoning into a real latency saving rather @@ -383,8 +385,6 @@ If you find edge0 useful, please cite our technical report: ## Contact Us - - Community and support channels are coming soon — this section will list the official ways to reach us: diff --git a/README_zh.md b/README_zh.md index 6b76609..da25752 100644 --- a/README_zh.md +++ b/README_zh.md @@ -22,11 +22,11 @@ ## 新闻 -- **【2026-09-30】** 我们发布了**三端推理引擎**，适配 iOS、macOS、Android - 平台，让用户在不同架构、不同平台上都有最佳的推理体验。三端源码已 - 开源至本仓库（[`ios/`](ios/) · [`macos/`](macos/) · - [`android/`](android/)），更多细节见各目录 README；**统一推理框架** - 将于 **2026 年 10 月底**发布，详见[路线图](#路线图)。 +- **【2026-09-30】** 我们发布了**四端推理引擎**，适配 iOS、macOS、 + Android、Windows 平台，让用户在不同架构、不同平台上都有最佳的推理 + 体验。四端源码已开源至本仓库（[`ios/`](ios/) · [`macos/`](macos/) · + [`android/`](android/) · [`windows/`](windows/)），更多细节见各目录 + README；**统一推理框架**将于 **2026 Q4** 发布，详见[路线图](#路线图)。 - **【2026-09-20】** `edge0 demo / chat / serve` 新增 `--prefill-ondemand` （[#112](https://github.com/Edge0-AI/Edge0/pull/112)）。 - **【2026-09-16】** 技术报告上线 arXiv：[The Other Half of the Memory @@ -64,7 +64,7 @@ MoE 模型跑在消费级硬件上——峰值内存由**激活**专家集而非 | **macOS** 桌面 App 与 CLI | [`macos/`](macos/) | Rust | ✅ 已开源（2026-09-30） | | **iOS** App | [`ios/`](ios/) | Swift + MLX Swift | ✅ 已开源（2026-09-30） | | **Android** App 与引擎 | [`android/`](android/) | Kotlin + 原生引擎 | ✅ 已开源（2026-09-30） | -| **Windows** | — | — | 🗺️ 路线图中 | +| **Windows** App 与引擎 | [`windows/`](windows/) | C++ + Vulkan | ✅ 已开源（2026-09-30） | ### 模型 @@ -273,9 +273,9 @@ engine.close() # 释放 mmap / 专家缓存 - [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) - 技术报告：[The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063)（[PDF](paper/main.pdf)） -### macOS / iOS / Android +### macOS / iOS / Android / Windows -三端推理引擎已开源在本仓库，更多细节见各目录 README： +四端推理引擎已开源在本仓库，更多细节见各目录 README： - **macOS**：本地 CLI / daemon / 桌面 App（Rust）—— 详见 [`macos/README.md`](macos/README.md) @@ -283,23 +283,26 @@ engine.close() # 释放 mmap / 专家缓存 [`ios/README.md`](ios/README.md) - **Android**：端侧 App + 原生引擎（Kotlin）—— 详见 [`android/README.md`](android/README.md) +- **Windows**：桌面 App + 原生引擎（C++ + Vulkan）—— 详见 + [`windows/README.md`](windows/README.md) **统一推理框架**（接入层统一，runtime 自动适配 iOS / macOS / Android / -Windows / Python）将于 **2026 年 10 月底**发布，见[路线图](#路线图)。 +Windows / Python）将于 **2026 Q4** 发布，见[路线图](#路线图)。 ## 路线图 +### 2026 Q4 + **平台与系统** -- **【2026 年 10 月底】edge0 统一推理框架开源** —— 发布统一推理框架： - **接入层统一**为一套 API（chat / serve / 端侧共用），**runtime 自动 - 适配不同硬件平台**——iOS、macOS、Android、Windows、Python；在本仓库 - 已开源的三端引擎（`ios/` · `macos/` · `android/`）之上统一收编。 -- **Windows** 平台支持（经统一框架）。 +- **edge0 统一推理框架开源** —— 发布统一推理框架：**接入层统一**为一套 + API（chat / serve / 端侧共用），**runtime 自动适配不同硬件平台**—— + iOS、macOS、Android、Windows、Python；在本仓库已开源的平台引擎 + （`ios/` · `macos/` · `android/` · `windows/`）之上统一收编。 - **CUDA 后端**（Python 框架）—— 插槽已预留在 `python/src/edge0/backends/cuda/`，核心代码零改动。 -**模型与算法 —— 2026 Q4** +**模型与算法** Q4 推进两条主线：把下一代架构引入框架，以及把潜在推理（latent thinking）变成真正的延迟收益、而不只是算力节省。 @@ -364,8 +367,6 @@ examples/demo.py # 最小 API walkthrough（edge0 demo 的等价代 ## 联系我们 - - 社区与支持渠道即将上线，本节将列出联系我们的官方方式： - **邮箱**：samuel@edge0.ai diff --git a/android/README.md b/android/README.md index 9be16bd..bb891b6 100644 --- a/android/README.md +++ b/android/README.md @@ -10,7 +10,7 @@ Recover-LoRA, and prerouter routing prediction. The next milestone is the **unified inference framework** — one access layer, runtime auto-adapting to the hardware platform (iOS / macOS / -Android / Windows / Python) — targeted for the **end of October 2026**; +Android / Windows / Python) — targeted for **Q4 2026**; see the [roadmap](../README.md#roadmap). See the [repository README](../README.md) for the multi-platform picture. @@ -19,5 +19,5 @@ See the [repository README](../README.md) for the multi-platform picture. > ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 > 与使用文档。Android runtime：Kotlin App + 原生推理引擎的端侧推理。 -> **统一推理框架**将于 **2026 年 10 月底**发布，详见 +> **统一推理框架**将于 **2026 Q4** 发布，详见 > [根 README 路线图](../README_zh.md)。 diff --git a/ios/README.md b/ios/README.md index 7fc05a3..292de71 100644 --- a/ios/README.md +++ b/ios/README.md @@ -10,7 +10,7 @@ Recover-LoRA, and prerouter routing prediction. The next milestone is the **unified inference framework** — one access layer, runtime auto-adapting to the hardware platform (iOS / macOS / -Android / Windows / Python) — targeted for the **end of October 2026**; +Android / Windows / Python) — targeted for **Q4 2026**; see the [roadmap](../README.md#roadmap). See the [repository README](../README.md) for the multi-platform picture. @@ -19,5 +19,5 @@ See the [repository README](../README.md) for the multi-platform picture. > ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 > 与使用文档。iOS runtime：Swift / SwiftUI + MLX Swift 的 iPhone 端侧 -> 推理。**统一推理框架**将于 **2026 年 10 月底**发布，详见 +> 推理。**统一推理框架**将于 **2026 Q4** 发布，详见 > [根 README 路线图](../README_zh.md)。 diff --git a/macos/README.md b/macos/README.md index e4efb5c..fbfdf84 100644 --- a/macos/README.md +++ b/macos/README.md @@ -11,7 +11,7 @@ and serves them locally behind an OpenAI-compatible API. The next milestone is the **unified inference framework** — one access layer, runtime auto-adapting to the hardware platform (iOS / macOS / -Android / Windows / Python) — targeted for the **end of October 2026**; +Android / Windows / Python) — targeted for **Q4 2026**; see the [roadmap](../README.md#roadmap). See the [repository README](../README.md) for the multi-platform picture. @@ -20,5 +20,5 @@ See the [repository README](../README.md) for the multi-platform picture. > ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 > 与使用文档。macOS runtime：Rust + Tauri 的本地 CLI / daemon / 桌面 -> App，Ollama 式体验 + OpenAI 兼容 API。**统一推理框架**将于 **2026 年 -> 10 月底**发布，详见[根 README 路线图](../README_zh.md)。 +> App，Ollama 式体验 + OpenAI 兼容 API。**统一推理框架**将于 **2026 Q4** +> 发布，详见[根 README 路线图](../README_zh.md)。 diff --git a/windows/README.md b/windows/README.md new file mode 100644 index 0000000..4296800 --- /dev/null +++ b/windows/README.md @@ -0,0 +1,25 @@ +# edge0 — Windows runtime + +> ✅ Open-sourced on **2026-09-30**. This README will be updated with +> full build & usage docs together with the source drop. + +The Windows runtime is a native C++17 port of the edge0 streaming core +(CMake build) with a **Vulkan** compute backend: the `edge0` core +(streaming mmap / cache / layer / MoE) is mirrored from the upstream +Python tree, running the same recipe as the rest of the framework — SSD +expert offload, Recover-LoRA, and prerouter routing prediction. + +The next milestone is the **unified inference framework** — one access +layer, runtime auto-adapting to the hardware platform (iOS / macOS / +Android / Windows / Python) — targeted for **Q4 2026**; see the +[roadmap](../README.md#roadmap). + +See the [repository README](../README.md) for the multi-platform picture. + +--- + +> ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 +> 与使用文档。Windows runtime：edge0 流式核心的原生 C++17 移植（CMake +> 构建）+ **Vulkan** 计算后端，`edge0` 核心（streaming mmap / cache / +> layer / MoE）镜像自上游 Python 树，沿用同一套配方。**统一推理框架** +> 将于 **2026 Q4** 发布，详见[根 README 路线图](../README_zh.md)。 From f93251a21b674fa4d9229e9874c71077d277303c Mon Sep 17 00:00:00 2001 From: Bin Date: Wed, 30 Sep 2026 14:55:45 +0800 Subject: [PATCH 08/13] add ios version --- ios/Config/App.xcconfig | 1 + ios/Config/Local.xcconfig.example | 3 + ios/Edge0PhoneProbe.xcodeproj/project.pbxproj | 259 +++++ .../contents.xcworkspacedata | 7 + .../xcshareddata/swiftpm/Package.resolved | 33 + .../xcschemes/Edge0Phone.xcscheme | 76 ++ ios/Edge0PhoneProbe/ContentView.swift | 552 +++++++++ .../Edge0PhoneProbe.entitlements | 6 + ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift | 10 + ios/Edge0PhoneProbe/Info.plist | 35 + ios/Models/.gitkeep | 0 ios/Package.resolved | 33 + ios/Package.swift | 33 + ios/README.md | 120 +- .../Edge0Core/Edge0Configuration8B.swift | 99 ++ ios/Sources/Edge0Core/Edge0Tokenizer8B.swift | 186 +++ ios/Sources/Edge0Core/ExpertTensorStore.swift | 60 + .../Edge0Core/GroupedExpertRouter.swift | 58 + ios/Sources/Edge0Core/MappedFile.swift | 123 ++ ios/Sources/Edge0Core/SafetensorsIndex.swift | 152 +++ .../Edge0MLX/Edge035B/ChatTemplate.swift | 122 ++ .../Edge035B/CoreMLVocabularyHead.swift | 60 + .../Edge035B/Edge0ChatEngine35B.swift | 284 +++++ .../Edge0MLX/Edge035B/Edge0Model35B.swift | 1012 +++++++++++++++++ .../Edge0MLX/Edge035B/ExpertLoader.swift | 241 ++++ .../Edge0MLX/Edge035B/PrefixCache.swift | 192 ++++ .../Edge0MLX/Edge035B/PregateHeads.swift | 220 ++++ .../Edge0MLX/Edge035B/RepackedExperts.swift | 162 +++ .../Edge0MLX/Edge035B/Safetensors.swift | 264 +++++ .../Edge0MLX/Edge035B/StepProfile.swift | 148 +++ .../Edge0MLX/Edge035B/ThinkingBlock.swift | 57 + ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift | 267 +++++ ios/Sources/Edge0MLX/Edge0ChatEngine.swift | 249 ++++ .../Edge0MLX/Edge0GroupedRouter8B.swift | 53 + ios/Sources/Edge0MLX/LingPrerouter.swift | 108 ++ .../Edge0MLX/MLXExpertTensorStore.swift | 56 + ios/Sources/Edge0MLX/OfficialGatedDelta.swift | 114 ++ ios/Sources/Edge0MLX/ParallelLoRA.swift | 62 + ios/Sources/Edge0MLX/QuantizedExpert.swift | 227 ++++ .../Edge0MLX/StreamingEdge0Model8B.swift | 295 +++++ .../Edge0MLX/StreamingKDADecoderBlock.swift | 285 +++++ .../Edge0MLX/StreamingMLADecoderBlock.swift | 278 +++++ ios/Sources/Edge0MLX/StreamingMoE.swift | 475 ++++++++ .../Edge0Tokenizer8BTests.swift | 12 + .../Edge0CoreTests/RouterReferenceTests.swift | 19 + .../SafetensorsIndexTests.swift | 32 + .../Edge0MLXTests/Edge035BChatTests.swift | 45 + ios/Tests/Edge0MLXTests/Fixtures/config.json | 105 ++ ios/Tests/Edge0MLXTests/M1Tests.swift | 161 +++ ios/scripts/test.sh | 17 + ios/tools/convert_pregate.py | 108 ++ ios/tools/convert_tokenizer.py | 110 ++ ios/tools/repack_experts.py | 354 ++++++ 53 files changed, 8023 insertions(+), 17 deletions(-) create mode 100644 ios/Config/App.xcconfig create mode 100644 ios/Config/Local.xcconfig.example create mode 100644 ios/Edge0PhoneProbe.xcodeproj/project.pbxproj create mode 100644 ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata create mode 100644 ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved create mode 100644 ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme create mode 100644 ios/Edge0PhoneProbe/ContentView.swift create mode 100644 ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements create mode 100644 ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift create mode 100644 ios/Edge0PhoneProbe/Info.plist create mode 100644 ios/Models/.gitkeep create mode 100644 ios/Package.resolved create mode 100644 ios/Package.swift create mode 100644 ios/Sources/Edge0Core/Edge0Configuration8B.swift create mode 100644 ios/Sources/Edge0Core/Edge0Tokenizer8B.swift create mode 100644 ios/Sources/Edge0Core/ExpertTensorStore.swift create mode 100644 ios/Sources/Edge0Core/GroupedExpertRouter.swift create mode 100644 ios/Sources/Edge0Core/MappedFile.swift create mode 100644 ios/Sources/Edge0Core/SafetensorsIndex.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/Safetensors.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/StepProfile.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift create mode 100644 ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift create mode 100644 ios/Sources/Edge0MLX/Edge0ChatEngine.swift create mode 100644 ios/Sources/Edge0MLX/Edge0GroupedRouter8B.swift create mode 100644 ios/Sources/Edge0MLX/LingPrerouter.swift create mode 100644 ios/Sources/Edge0MLX/MLXExpertTensorStore.swift create mode 100644 ios/Sources/Edge0MLX/OfficialGatedDelta.swift create mode 100644 ios/Sources/Edge0MLX/ParallelLoRA.swift create mode 100644 ios/Sources/Edge0MLX/QuantizedExpert.swift create mode 100644 ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift create mode 100644 ios/Sources/Edge0MLX/StreamingKDADecoderBlock.swift create mode 100644 ios/Sources/Edge0MLX/StreamingMLADecoderBlock.swift create mode 100644 ios/Sources/Edge0MLX/StreamingMoE.swift create mode 100644 ios/Tests/Edge0CoreTests/Edge0Tokenizer8BTests.swift create mode 100644 ios/Tests/Edge0CoreTests/RouterReferenceTests.swift create mode 100644 ios/Tests/Edge0CoreTests/SafetensorsIndexTests.swift create mode 100644 ios/Tests/Edge0MLXTests/Edge035BChatTests.swift create mode 100644 ios/Tests/Edge0MLXTests/Fixtures/config.json create mode 100644 ios/Tests/Edge0MLXTests/M1Tests.swift create mode 100755 ios/scripts/test.sh create mode 100755 ios/tools/convert_pregate.py create mode 100755 ios/tools/convert_tokenizer.py create mode 100755 ios/tools/repack_experts.py diff --git a/ios/Config/App.xcconfig b/ios/Config/App.xcconfig new file mode 100644 index 0000000..234f81e --- /dev/null +++ b/ios/Config/App.xcconfig @@ -0,0 +1 @@ +#include? "Local.xcconfig" diff --git a/ios/Config/Local.xcconfig.example b/ios/Config/Local.xcconfig.example new file mode 100644 index 0000000..80c6e5d --- /dev/null +++ b/ios/Config/Local.xcconfig.example @@ -0,0 +1,3 @@ +// Copy this file to Config/Local.xcconfig and set your Apple Developer Team ID. +// Config/Local.xcconfig is excluded by .gitignore. +DEVELOPMENT_TEAM = diff --git a/ios/Edge0PhoneProbe.xcodeproj/project.pbxproj b/ios/Edge0PhoneProbe.xcodeproj/project.pbxproj new file mode 100644 index 0000000..4bdcaac --- /dev/null +++ b/ios/Edge0PhoneProbe.xcodeproj/project.pbxproj @@ -0,0 +1,259 @@ +// !$*UTF8*$! +{ + archiveVersion = 1; + classes = { + }; + objectVersion = 71; + objects = { + +/* Begin PBXBuildFile section */ + A00000000000000000000001 /* Edge0PhoneProbeApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = A00000000000000000000011 /* Edge0PhoneProbeApp.swift */; }; + A00000000000000000000002 /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = A00000000000000000000012 /* ContentView.swift */; }; + A00000000000000000000004 /* Edge0Core in Frameworks */ = {isa = PBXBuildFile; productRef = A00000000000000000000071 /* Edge0Core */; }; + A00000000000000000000005 /* Edge0MLX in Frameworks */ = {isa = PBXBuildFile; productRef = A00000000000000000000072 /* Edge0MLX */; }; + A00000000000000000000006 /* Models in Resources */ = {isa = PBXBuildFile; fileRef = A00000000000000000000014 /* Models */; }; +/* End PBXBuildFile section */ + +/* Begin PBXFileReference section */ + A00000000000000000000003 /* Edge0Phone.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = Edge0Phone.app; sourceTree = BUILT_PRODUCTS_DIR; }; + A00000000000000000000011 /* Edge0PhoneProbeApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Edge0PhoneProbeApp.swift; sourceTree = ""; }; + A00000000000000000000012 /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = ""; }; + A00000000000000000000013 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = ""; }; + A00000000000000000000014 /* Models */ = {isa = PBXFileReference; lastKnownFileType = folder; path = Models; sourceTree = ""; }; + A00000000000000000000015 /* App.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = App.xcconfig; sourceTree = ""; }; +/* End PBXFileReference section */ + +/* Begin PBXFrameworksBuildPhase section */ + A00000000000000000000021 /* Frameworks */ = { + isa = PBXFrameworksBuildPhase; + buildActionMask = 2147483647; + files = ( + A00000000000000000000004 /* Edge0Core in Frameworks */, + A00000000000000000000005 /* Edge0MLX in Frameworks */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXFrameworksBuildPhase section */ + +/* Begin PBXGroup section */ + A00000000000000000000022 = { + isa = PBXGroup; + children = ( + A00000000000000000000033 /* Config */, + A00000000000000000000014 /* Models */, + A00000000000000000000031 /* Edge0PhoneProbe */, + A00000000000000000000032 /* Products */, + ); + sourceTree = ""; + }; + A00000000000000000000031 /* Edge0PhoneProbe */ = { + isa = PBXGroup; + children = ( + A00000000000000000000011 /* Edge0PhoneProbeApp.swift */, + A00000000000000000000012 /* ContentView.swift */, + A00000000000000000000013 /* Info.plist */, + ); + path = Edge0PhoneProbe; + sourceTree = ""; + }; + A00000000000000000000032 /* Products */ = { + isa = PBXGroup; + children = ( + A00000000000000000000003 /* Edge0Phone.app */, + ); + name = Products; + sourceTree = ""; + }; + A00000000000000000000033 /* Config */ = { + isa = PBXGroup; + children = ( + A00000000000000000000015 /* App.xcconfig */, + ); + path = Config; + sourceTree = ""; + }; +/* End PBXGroup section */ + +/* Begin PBXNativeTarget section */ + A00000000000000000000041 /* Edge0Phone */ = { + isa = PBXNativeTarget; + buildConfigurationList = A00000000000000000000051 /* Build configuration list for PBXNativeTarget "Edge0Phone" */; + buildPhases = ( + A00000000000000000000061 /* Sources */, + A00000000000000000000021 /* Frameworks */, + A00000000000000000000062 /* Resources */, + ); + buildRules = ( + ); + dependencies = ( + ); + name = Edge0Phone; + packageProductDependencies = ( + A00000000000000000000071 /* Edge0Core */, + A00000000000000000000072 /* Edge0MLX */, + ); + productName = Edge0Phone; + productReference = A00000000000000000000003 /* Edge0Phone.app */; + productType = "com.apple.product-type.application"; + }; +/* End PBXNativeTarget section */ + +/* Begin PBXProject section */ + A00000000000000000000042 /* Project object */ = { + isa = PBXProject; + attributes = { + BuildIndependentTargetsInParallel = 1; + LastSwiftUpdateCheck = 2610; + LastUpgradeCheck = 2610; + TargetAttributes = { + A00000000000000000000041 = { + CreatedOnToolsVersion = 26.6; + }; + }; + }; + buildConfigurationList = A00000000000000000000052 /* Build configuration list for PBXProject "Edge0PhoneProbe" */; + compatibilityVersion = "Xcode 16.0"; + developmentRegion = en; + hasScannedForEncodings = 0; + knownRegions = ( + en, + Base, + ); + mainGroup = A00000000000000000000022; + packageReferences = ( + A00000000000000000000081 /* XCLocalSwiftPackageReference "." */, + ); + productRefGroup = A00000000000000000000032 /* Products */; + projectDirPath = ""; + projectRoot = ""; + targets = ( + A00000000000000000000041 /* Edge0Phone */, + ); + }; +/* End PBXProject section */ + +/* Begin PBXResourcesBuildPhase section */ + A00000000000000000000062 /* Resources */ = { + isa = PBXResourcesBuildPhase; + buildActionMask = 2147483647; + files = ( + A00000000000000000000006 /* Models in Resources */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXResourcesBuildPhase section */ + +/* Begin PBXSourcesBuildPhase section */ + A00000000000000000000061 /* Sources */ = { + isa = PBXSourcesBuildPhase; + buildActionMask = 2147483647; + files = ( + A00000000000000000000001 /* Edge0PhoneProbeApp.swift in Sources */, + A00000000000000000000002 /* ContentView.swift in Sources */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXSourcesBuildPhase section */ + +/* Begin XCBuildConfiguration section */ + A00000000000000000000091 /* Debug */ = { + isa = XCBuildConfiguration; + baseConfigurationReference = A00000000000000000000015 /* App.xcconfig */; + buildSettings = { + CODE_SIGN_ENTITLEMENTS = Edge0PhoneProbe/Edge0PhoneProbe.entitlements; + CODE_SIGN_STYLE = Automatic; + DEVELOPMENT_TEAM = X756QY2CTK; + GENERATE_INFOPLIST_FILE = NO; + INFOPLIST_FILE = Edge0PhoneProbe/Info.plist; + IPHONEOS_DEPLOYMENT_TARGET = 17.0; + PRODUCT_BUNDLE_IDENTIFIER = com.example.edge0phone; + PRODUCT_NAME = "$(TARGET_NAME)"; + SDKROOT = iphoneos; + SUPPORTED_PLATFORMS = "iphoneos iphonesimulator"; + SWIFT_VERSION = 6.0; + TARGETED_DEVICE_FAMILY = 1; + }; + name = Debug; + }; + A00000000000000000000092 /* Release */ = { + isa = XCBuildConfiguration; + baseConfigurationReference = A00000000000000000000015 /* App.xcconfig */; + buildSettings = { + CODE_SIGN_ENTITLEMENTS = Edge0PhoneProbe/Edge0PhoneProbe.entitlements; + CODE_SIGN_STYLE = Automatic; + DEVELOPMENT_TEAM = X756QY2CTK; + GENERATE_INFOPLIST_FILE = NO; + INFOPLIST_FILE = Edge0PhoneProbe/Info.plist; + IPHONEOS_DEPLOYMENT_TARGET = 17.0; + PRODUCT_BUNDLE_IDENTIFIER = com.example.edge0phone; + PRODUCT_NAME = "$(TARGET_NAME)"; + SDKROOT = iphoneos; + SUPPORTED_PLATFORMS = "iphoneos iphonesimulator"; + SWIFT_VERSION = 6.0; + TARGETED_DEVICE_FAMILY = 1; + }; + name = Release; + }; + A00000000000000000000093 /* Debug */ = { + isa = XCBuildConfiguration; + buildSettings = { + ALWAYS_SEARCH_USER_PATHS = NO; + CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; + CLANG_ENABLE_MODULES = YES; + }; + name = Debug; + }; + A00000000000000000000094 /* Release */ = { + isa = XCBuildConfiguration; + buildSettings = { + ALWAYS_SEARCH_USER_PATHS = NO; + CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; + CLANG_ENABLE_MODULES = YES; + }; + name = Release; + }; +/* End XCBuildConfiguration section */ + +/* Begin XCConfigurationList section */ + A00000000000000000000051 /* Build configuration list for PBXNativeTarget "Edge0Phone" */ = { + isa = XCConfigurationList; + buildConfigurations = ( + A00000000000000000000091 /* Debug */, + A00000000000000000000092 /* Release */, + ); + defaultConfigurationIsVisible = 0; + defaultConfigurationName = Release; + }; + A00000000000000000000052 /* Build configuration list for PBXProject "Edge0PhoneProbe" */ = { + isa = XCConfigurationList; + buildConfigurations = ( + A00000000000000000000093 /* Debug */, + A00000000000000000000094 /* Release */, + ); + defaultConfigurationIsVisible = 0; + defaultConfigurationName = Release; + }; +/* End XCConfigurationList section */ + +/* Begin XCLocalSwiftPackageReference section */ + A00000000000000000000081 /* XCLocalSwiftPackageReference "." */ = { + isa = XCLocalSwiftPackageReference; + relativePath = .; + }; +/* End XCLocalSwiftPackageReference section */ + +/* Begin XCSwiftPackageProductDependency section */ + A00000000000000000000071 /* Edge0Core */ = { + isa = XCSwiftPackageProductDependency; + package = A00000000000000000000081 /* XCLocalSwiftPackageReference "." */; + productName = Edge0Core; + }; + A00000000000000000000072 /* Edge0MLX */ = { + isa = XCSwiftPackageProductDependency; + package = A00000000000000000000081 /* XCLocalSwiftPackageReference "." */; + productName = Edge0MLX; + }; +/* End XCSwiftPackageProductDependency section */ + }; + rootObject = A00000000000000000000042 /* Project object */; +} diff --git a/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata new file mode 100644 index 0000000..919434a --- /dev/null +++ b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata @@ -0,0 +1,7 @@ + + + + + diff --git a/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved new file mode 100644 index 0000000..e1dd047 --- /dev/null +++ b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved @@ -0,0 +1,33 @@ +{ + "originHash" : "bb8f886c2113b757a4f1caf03c3e9422df8f501d7a701d5dfdb178ee6043c4eb", + "pins" : [ + { + "identity" : "mlx-swift", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ml-explore/mlx-swift", + "state" : { + "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021", + "version" : "0.31.6" + } + }, + { + "identity" : "swift-argument-parser", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-argument-parser", + "state" : { + "revision" : "6a52f3251125d74daf04fcbd5e6f08a75d074382", + "version" : "1.8.2" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics", + "state" : { + "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", + "version" : "1.1.1" + } + } + ], + "version" : 3 +} diff --git a/ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme b/ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme new file mode 100644 index 0000000..0136b32 --- /dev/null +++ b/ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme @@ -0,0 +1,76 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/ios/Edge0PhoneProbe/ContentView.swift b/ios/Edge0PhoneProbe/ContentView.swift new file mode 100644 index 0000000..4e33ad3 --- /dev/null +++ b/ios/Edge0PhoneProbe/ContentView.swift @@ -0,0 +1,552 @@ +import Edge0MLX +import Darwin +import Foundation +import SwiftUI + +private struct ChatMessage: Identifiable, Sendable { + enum Role: Sendable { case user, assistant } + let id: UUID + let role: Role + var text: String + var detail: String? + + init(id: UUID = UUID(), role: Role, text: String, detail: String?) { + self.id = id; self.role = role; self.text = text; self.detail = detail + } +} + +private enum LocalModel: String, CaseIterable, Identifiable, Sendable { + case edge8 + case edge35b + + var id: String { rawValue } + var title: String { + switch self { + case .edge8: "Edge0 8B" + case .edge35b: "Edge0 35B" + } + } + var folderName: String { + self == .edge8 ? "Edge0-8B-A1B-preview" : Edge0ChatEngine35B.modelFolderName + } +} + +private actor ChatRuntime { + private enum LoadedEngine { + case edge8(Edge0ChatEngine) + case edge35b(Edge0ChatEngine35B) + } + private var engine: LoadedEngine? + + func load(_ model: LocalModel, modelURL: URL) throws { + // Do not keep both models alive while switching; 35B is memory constrained. + engine = nil + switch model { + case .edge8: + engine = .edge8(try Edge0ChatEngine(modelURL: modelURL)) + case .edge35b: + engine = .edge35b(try Edge0ChatEngine35B(modelURL: modelURL)) + } + } + + func reply(to text: String, + maxTokens: Int? = nil, + thinking: Bool, + onText: @escaping @Sendable (String) -> Void) async throws -> Edge0GenerationResult { + let shouldContinue = { @Sendable in + !withUnsafeCurrentTask { $0?.isCancelled ?? false } + } + switch engine { + case .edge8(let engine): + if let maxTokens { + return try engine.reply( + to: text, maxTokens: maxTokens, thinking: thinking, + onText: onText, shouldContinue: shouldContinue) + } + return try engine.reply( + to: text, thinking: thinking, onText: onText, shouldContinue: shouldContinue) + case .edge35b(let engine): + if let maxTokens { + return try await engine.reply( + to: text, maxTokens: maxTokens, thinking: thinking, + onText: onText, shouldContinue: shouldContinue) + } + return try await engine.reply( + to: text, thinking: thinking, onText: onText, shouldContinue: shouldContinue) + case nil: + throw RuntimeFailure.modelNotLoaded + } + } + + func reset() { + switch engine { + case .edge8(let engine): engine.reset() + case .edge35b(let engine): engine.reset() + case nil: break + } + } + + func unload() { engine = nil } + + private enum RuntimeFailure: LocalizedError { + case modelNotLoaded + var errorDescription: String? { "Model is not loaded" } + } +} + +@MainActor +private final class ChatViewModel: ObservableObject { + enum Phase: Equatable { case choosing, ready, preparing, generating, failed(String) } + + @Published var input = "" + @Published private(set) var messages: [ChatMessage] = [] + @Published private(set) var phase: Phase = .choosing + @Published private(set) var selectedModel: LocalModel? + + private let runtime = ChatRuntime() + private var generationTask: Task? + + var isRunning: Bool { phase == .preparing || phase == .generating } + /// Off unless the process is launched with `--thinking`. `--no-thinking` forces it off. + private let thinkingEnabled: Bool = { + let args = CommandLine.arguments + if args.contains("--no-thinking") { return false } + return args.contains("--thinking") + }() + var statusText: String? { + switch phase { + case .choosing: nil + case .ready: nil + case .preparing: "Loading \(selectedModel?.title ?? "local model")…" + case .generating: "\(selectedModel?.title ?? "Model") is generating…" + case .failed(let message): message + } + } + + func isInstalled(_ model: LocalModel) -> Bool { + let url = resolveModelURL(model) + switch model { + case .edge8: + return FileManager.default.fileExists( + atPath: url.appendingPathComponent("model.safetensors").path) + case .edge35b: + return FileManager.default.fileExists( + atPath: url.appendingPathComponent("model.safetensors.index.json").path) + && FileManager.default.fileExists( + atPath: url.appendingPathComponent("tokenizer.bin").path) + && FileManager.default.fileExists( + atPath: url.appendingPathComponent("experts-L39.bin").path) + } + } + + func choose(_ model: LocalModel) { + guard !isRunning else { return } + selectedModel = model + messages.removeAll() + input = "" + phase = .preparing + let url = resolveModelURL(model) + Task { + do { + try await runtime.load(model, modelURL: url) + if phase == .preparing { phase = .ready } + } catch { + phase = .failed(Self.friendlyMessage(for: error)) + } + } + } + + func chooseAnotherModel() { + guard !isRunning else { return } + messages.removeAll() + input = "" + selectedModel = nil + phase = .choosing + Task { await runtime.unload() } + } + + func useSuggestion(_ text: String, maxTokens: Int? = nil) { + input = text + send(maxTokens: maxTokens) + } + + func send(maxTokens: Int? = nil) { + let text = input.trimmingCharacters(in: .whitespacesAndNewlines) + guard !text.isEmpty, !isRunning else { return } + input = "" + messages.append(ChatMessage(role: .user, text: text, detail: nil)) + let assistantID = UUID() + messages.append(ChatMessage(id: assistantID, role: .assistant, + text: "", detail: nil)) + phase = .generating + generationTask = Task { + do { + let result = try await runtime.reply( + to: text, + maxTokens: maxTokens, + thinking: thinkingEnabled, + onText: { [weak self] partial in + Task { @MainActor in + guard let self, + let index = self.messages.firstIndex(where: { $0.id == assistantID }) + else { return } + self.messages[index].text = partial + } + }) + let detail = String( + format: "%d tokens · TTFT %.1fs · prefill %.1f tok/s · decode %.2f tok/s · peak %.0f MB", + result.generatedTokenCount, result.timeToFirstTokenSeconds, + result.prefillTokensPerSecond, + result.decodeTokensPerSecond, + Double(result.peakMemoryBytes) / 1_048_576 + ) + if let index = messages.firstIndex(where: { $0.id == assistantID }) { + messages[index].text = result.text.isEmpty ? "(The model returned an empty response)" : result.text + messages[index].detail = detail + } + print("[generation] \(detail)") + if CommandLine.arguments.contains("--bench-second"), + CommandLine.arguments.contains("--bench-35b") { + print("[bench] second turn") + let second = try await runtime.reply( + to: "Summarize it in one sentence.", + maxTokens: 64, + thinking: thinkingEnabled, + onText: { _ in }) + let secondDetail = String( + format: "%d tokens · TTFT %.1fs · prefill %.1f tok/s · decode %.2f tok/s · peak %.0f MB", + second.generatedTokenCount, second.timeToFirstTokenSeconds, + second.prefillTokensPerSecond, + second.decodeTokensPerSecond, + Double(second.peakMemoryBytes) / 1_048_576 + ) + print("[generation turn2] \(secondDetail)") + } + if CommandLine.arguments.contains("--bench-8b") + || CommandLine.arguments.contains("--bench-35b") { + var info = task_vm_info_data_t() + var count = mach_msg_type_number_t( + MemoryLayout.stride / MemoryLayout.stride) + let status = withUnsafeMutablePointer(to: &info) { + $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { + task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count) + } + } + print("[bench memory] endFootprintMiB=\(status == KERN_SUCCESS ? Double(info.phys_footprint) / 1_048_576 : -1) thermal=\(ProcessInfo.processInfo.thermalState.rawValue)") + print("[bench timing] ttftSeconds=\(result.timeToFirstTokenSeconds) prefillSeconds=\(result.prefillSeconds) prefillTokPerSec=\(result.prefillTokensPerSecond) decodeTokPerSec=\(result.decodeTokensPerSecond)") + } + phase = .ready + if CommandLine.arguments.contains("--bench-8b") + || CommandLine.arguments.contains("--bench-35b") + || CommandLine.arguments.contains("--smoke-35b") { + fflush(nil) + exit(0) + } + } catch is CancellationError { + phase = .ready + } catch { + phase = .failed(Self.friendlyMessage(for: error)) + if CommandLine.arguments.contains("--bench-8b") + || CommandLine.arguments.contains("--bench-35b") + || CommandLine.arguments.contains("--smoke-35b") { + fflush(stdout) + exit(1) + } + } + generationTask = nil + } + } + + func stop() { generationTask?.cancel() } + + func newConversation() { + guard !isRunning, selectedModel != nil else { return } + messages.removeAll() + input = "" + phase = .ready + Task { await runtime.reset() } + } + + private func resolveModelURL(_ model: LocalModel) -> URL { + let models = Bundle.main.resourceURL? + .appendingPathComponent("Models", isDirectory: true) + ?? Bundle.main.bundleURL.appendingPathComponent("Models", isDirectory: true) + return models.appendingPathComponent(model.folderName, isDirectory: true) + } + + private static func friendlyMessage(for error: Error) -> String { + let description = error.localizedDescription + if description.contains("doesn’t exist") || description.contains("No such file") + || description.contains("incomplete") { + return "Model not included in this build" + } + return "Run failed: \(description)" + } +} + +struct ContentView: View { + @StateObject private var chat = ChatViewModel() + @FocusState private var inputFocused: Bool + + private let suggestions = [ + "What is artificial intelligence?", + "How do LLMs understand context?", + "Tell me about renewable energy", + "Write a short poem about autumn", + ] + + var body: some View { + ZStack { + Color.black.ignoresSafeArea() + if chat.selectedModel == nil { + modelPicker + } else { + VStack(spacing: 0) { + header + conversation + } + } + } + .preferredColorScheme(.dark) + .safeAreaInset(edge: .bottom, spacing: 0) { + if chat.selectedModel != nil { composer } + } + .task { + // Keeps the normal launch picker intact while allowing a repeatable + // on-device smoke/benchmark run from `devicectl`. + let smoke = CommandLine.arguments.contains("--smoke-35b") + let bench8 = CommandLine.arguments.contains("--bench-8b") + let bench35 = CommandLine.arguments.contains("--bench-35b") + let benchmark = bench8 || bench35 + let target: LocalModel = + bench8 || CommandLine.arguments.contains("--model=8b") ? .edge8 : .edge35b + if chat.selectedModel == nil, + (CommandLine.arguments.contains("--model=35b") + || CommandLine.arguments.contains("--model=8b") + || smoke || benchmark), + chat.isInstalled(target) { + print("[bench] loading \(target.title)") + print("[bench] thermalAtLaunch=\(ProcessInfo.processInfo.thermalState.rawValue)") + chat.choose(target) + if smoke || benchmark { + while chat.phase == .preparing { + try? await Task.sleep(for: .milliseconds(100)) + } + if case .failed(let message) = chat.phase { + print("[bench] load failed: \(message)") + } else if chat.phase == .ready { + let long = CommandLine.arguments.contains("--bench-long") + let sentence = "Streaming inference overlaps disk bandwidth with compute. Each mixture-of-experts layer activates only a few experts." + let prompt = smoke + ? "Reply with this only: 35B is running" + : (long + ? Array(repeating: sentence, count: 80).joined() + : "What is artificial intelligence?") + chat.useSuggestion(prompt, maxTokens: long ? 8 : (benchmark ? 64 : nil)) + } + } + } + } + } + + private var modelPicker: some View { + ScrollView { + VStack(alignment: .leading, spacing: 22) { + Spacer(minLength: 54) + Text("Choose a local model") + .font(.system(size: 34, weight: .bold, design: .rounded)) + + ForEach(LocalModel.allCases) { model in + let installed = chat.isInstalled(model) + Button { chat.choose(model) } label: { + HStack(spacing: 16) { + Image(systemName: model == .edge8 ? "bolt.fill" : "brain.head.profile") + .font(.system(size: 23, weight: .semibold)) + .frame(width: 48, height: 48) + .background(.white.opacity(0.09), in: RoundedRectangle(cornerRadius: 14)) + VStack(alignment: .leading, spacing: 5) { + Text(model.title).font(.headline) + if !installed { + Text("Not included") + .font(.caption2) + .foregroundStyle(.orange) + } + } + Spacer() + Image(systemName: "chevron.right") + .foregroundStyle(.secondary) + } + .padding(18) + .background(.white.opacity(0.07), in: RoundedRectangle(cornerRadius: 22)) + } + .buttonStyle(.plain) + .disabled(!installed) + .opacity(installed ? 1 : 0.62) + } + + Spacer(minLength: 40) + } + .padding(.horizontal, 22) + .frame(maxWidth: 620) + .frame(maxWidth: .infinity) + } + } + + private var header: some View { + HStack(spacing: 10) { + Text(chat.selectedModel?.title ?? "Edge0") + .font(.headline.weight(.semibold)) + Spacer() + Button(action: chat.chooseAnotherModel) { + Image(systemName: "arrow.left.arrow.right") + .font(.system(size: 15, weight: .semibold)) + .frame(width: 38, height: 38) + .background(.white.opacity(0.07), in: Circle()) + } + .foregroundStyle(.white) + .disabled(chat.isRunning) + .accessibilityLabel("Switch model") + Button(action: chat.newConversation) { + Image(systemName: "square.and.pencil") + .font(.system(size: 17, weight: .semibold)) + .frame(width: 38, height: 38) + .background(.white.opacity(0.07), in: Circle()) + } + .foregroundStyle(.white) + .disabled(chat.isRunning) + .accessibilityLabel("New chat") + } + .padding(.horizontal, 20) + .padding(.top, 8) + .padding(.bottom, 10) + } + + @ViewBuilder + private var conversation: some View { + if chat.messages.isEmpty { + welcome + } else { + ScrollViewReader { proxy in + ScrollView { + LazyVStack(spacing: 18) { + ForEach(chat.messages) { message in + messageBubble(message).id(message.id) + } + if let status = chat.statusText { statusRow(status).id("status") } + } + .padding(.horizontal, 18) + .padding(.top, 16) + .padding(.bottom, 24) + } + .scrollDismissesKeyboard(.interactively) + .onChange(of: chat.messages.count) { + withAnimation { proxy.scrollTo(chat.messages.last?.id, anchor: .bottom) } + } + .onChange(of: chat.phase) { + withAnimation { proxy.scrollTo("status", anchor: .bottom) } + } + } + } + } + + private var welcome: some View { + ScrollView { + VStack(spacing: 18) { + Spacer(minLength: 90) + Text("edge0") + .font(.system(size: 58, weight: .medium, design: .rounded)) + .tracking(-3) + Text("Private AI, running locally on your iPhone") + .font(.subheadline) + .foregroundStyle(.secondary) + .multilineTextAlignment(.center) + VStack(spacing: 10) { + ForEach(suggestions, id: \.self) { suggestion in + Button(suggestion) { chat.useSuggestion(suggestion) } + .font(.subheadline.weight(.medium)) + .foregroundStyle(.white.opacity(0.88)) + .padding(.horizontal, 16) + .padding(.vertical, 11) + .background(.white.opacity(0.07), in: Capsule()) + } + } + .padding(.top, 8) + if let status = chat.statusText { statusRow(status) } + Spacer(minLength: 120) + } + .frame(maxWidth: .infinity) + .padding(.horizontal, 24) + } + } + + private func messageBubble(_ message: ChatMessage) -> some View { + VStack(alignment: message.role == .user ? .trailing : .leading, spacing: 6) { + Text(message.text) + .font(.body) + .foregroundStyle(.white.opacity(0.94)) + .textSelection(.enabled) + .padding(.horizontal, 15) + .padding(.vertical, 12) + .background( + message.role == .user ? Color.white.opacity(0.14) : Color.white.opacity(0.07), + in: RoundedRectangle(cornerRadius: 19, style: .continuous) + ) + if let detail = message.detail { + Text(detail) + .font(.caption2.monospacedDigit()) + .foregroundStyle(.secondary) + .padding(.horizontal, 6) + } + } + .frame(maxWidth: .infinity, alignment: message.role == .user ? .trailing : .leading) + } + + private func statusRow(_ text: String) -> some View { + HStack(spacing: 10) { + if chat.isRunning { ProgressView().controlSize(.small) } + Text(text) + .font(.footnote) + .foregroundStyle(chat.isRunning ? Color.secondary : Color.red.opacity(0.9)) + } + .frame(maxWidth: .infinity, alignment: .leading) + .padding(.horizontal, 6) + } + + private var composer: some View { + HStack(alignment: .bottom, spacing: 10) { + TextField("Ask anything…", text: $chat.input, axis: .vertical) + .lineLimit(1...5) + .focused($inputFocused) + .submitLabel(.send) + .onSubmit { chat.send() } + .padding(.horizontal, 16) + .padding(.vertical, 13) + .background(.white.opacity(0.08), in: RoundedRectangle(cornerRadius: 22, style: .continuous)) + + Button { + if chat.isRunning, chat.phase == .generating { + chat.stop() + } else { + chat.send() + } + inputFocused = false + } label: { + Image(systemName: chat.phase == .generating ? "stop.fill" : "arrow.up") + .font(.system(size: 17, weight: .bold)) + .foregroundStyle(.black) + .frame(width: 45, height: 45) + .background(.white, in: Circle()) + } + .disabled(chat.phase == .preparing || + chat.phase == .choosing || + (chat.phase != .generating && + chat.input.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty)) + .opacity(chat.phase == .preparing ? 0.5 : 1) + .accessibilityLabel(chat.phase == .generating ? "Stop" : "Send") + } + .padding(.horizontal, 16) + .padding(.top, 10) + .padding(.bottom, 8) + .background(.ultraThinMaterial) + } +} diff --git a/ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements b/ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements new file mode 100644 index 0000000..6631ffa --- /dev/null +++ b/ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements @@ -0,0 +1,6 @@ + + + + + + diff --git a/ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift b/ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift new file mode 100644 index 0000000..0423c3d --- /dev/null +++ b/ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift @@ -0,0 +1,10 @@ +import SwiftUI + +@main +struct Edge0PhoneProbeApp: App { + var body: some Scene { + WindowGroup { + ContentView() + } + } +} diff --git a/ios/Edge0PhoneProbe/Info.plist b/ios/Edge0PhoneProbe/Info.plist new file mode 100644 index 0000000..91ac1d1 --- /dev/null +++ b/ios/Edge0PhoneProbe/Info.plist @@ -0,0 +1,35 @@ + + + + + CFBundleDevelopmentRegion + $(DEVELOPMENT_LANGUAGE) + CFBundleExecutable + $(EXECUTABLE_NAME) + CFBundleIdentifier + $(PRODUCT_BUNDLE_IDENTIFIER) + CFBundleInfoDictionaryVersion + 6.0 + CFBundleName + $(PRODUCT_NAME) + CFBundlePackageType + $(PRODUCT_BUNDLE_PACKAGE_TYPE) + CFBundleShortVersionString + 1.0 + CFBundleVersion + 1 + LSRequiresIPhoneOS + + LSSupportsOpeningDocumentsInPlace + + UIApplicationSceneManifest + + UIApplicationSupportsMultipleScenes + + + UIFileSharingEnabled + + UILaunchScreen + + + diff --git a/ios/Models/.gitkeep b/ios/Models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/ios/Package.resolved b/ios/Package.resolved new file mode 100644 index 0000000..66a2fa4 --- /dev/null +++ b/ios/Package.resolved @@ -0,0 +1,33 @@ +{ + "originHash" : "14c443ca8e8c4734ac842efe15d1c40af0856e85f2a4579f1ff7b75f93b2ccc7", + "pins" : [ + { + "identity" : "mlx-swift", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ml-explore/mlx-swift", + "state" : { + "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021", + "version" : "0.31.6" + } + }, + { + "identity" : "swift-argument-parser", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-argument-parser", + "state" : { + "revision" : "6a52f3251125d74daf04fcbd5e6f08a75d074382", + "version" : "1.8.2" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics", + "state" : { + "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", + "version" : "1.1.1" + } + } + ], + "version" : 3 +} diff --git a/ios/Package.swift b/ios/Package.swift new file mode 100644 index 0000000..4672abb --- /dev/null +++ b/ios/Package.swift @@ -0,0 +1,33 @@ +// swift-tools-version: 6.1 +import PackageDescription + +let package = Package( + name: "Edge0Phone", + platforms: [ + .macOS(.v14), + .iOS(.v17), + ], + products: [ + .library(name: "Edge0Core", targets: ["Edge0Core"]), + .library(name: "Edge0MLX", targets: ["Edge0MLX"]), + ], + dependencies: [ + .package( + url: "https://github.com/ml-explore/mlx-swift", + .upToNextMinor(from: "0.31.4") + ), + ], + targets: [ + .target(name: "Edge0Core"), + .target( + name: "Edge0MLX", + dependencies: [ + "Edge0Core", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXFast", package: "mlx-swift"), + ] + ), + .testTarget(name: "Edge0MLXTests", dependencies: ["Edge0MLX", "Edge0Core"], resources: [.copy("Fixtures")]), + .testTarget(name: "Edge0CoreTests", dependencies: ["Edge0Core"]), + ] +) diff --git a/ios/README.md b/ios/README.md index 292de71..bcde788 100644 --- a/ios/README.md +++ b/ios/README.md @@ -1,23 +1,109 @@ -# edge0 — iOS runtime +# Edge0Phone -> ✅ Open-sourced on **2026-09-30**. This README will be updated with -> full build & usage docs together with the source drop. +On-device chat for iPhone. The app can run Edge0 8B, Edge0 35B, or both, and only the model you select is loaded. Model weights are not in this folder. -The iOS runtime runs the edge0 models fully on-device on iPhone -(Swift / SwiftUI + MLX Swift): local streaming MoE inference with the -same recipe as the rest of the framework — SSD expert offload, -Recover-LoRA, and prerouter routing prediction. +## Quick Start -The next milestone is the **unified inference framework** — one access -layer, runtime auto-adapting to the hardware platform (iOS / macOS / -Android / Windows / Python) — targeted for **Q4 2026**; -see the [roadmap](../README.md#roadmap). +You need an Apple Silicon Mac, Xcode, and an iPhone on iOS 17 or later with Developer Mode on. Python 3.10 or later is only needed to download weights and to convert 35B. If Xcode reports a missing Metal compiler, run `xcodebuild -downloadComponent MetalToolchain`. -See the [repository README](../README.md) for the multi-platform picture. +The `Models/` directory is copied into the app at build time. Put a model there before you build if you want that button to be selectable. An empty `Models/` directory still builds; the missing model shows **Not included**. ---- +### Download a model -> ✅ 源码已于 **2026-09-30** 开源，本 README 将随源码落库补全完整的构建 -> 与使用文档。iOS runtime：Swift / SwiftUI + MLX Swift 的 iPhone 端侧 -> 推理。**统一推理框架**将于 **2026 Q4** 发布，详见 -> [根 README 路线图](../README_zh.md)。 +```sh +python3 -m venv .venv-tools +.venv-tools/bin/python -m pip install --upgrade pip huggingface_hub numpy +``` + +8B can be packed as downloaded. It is about 5 GB, so the installed app is about that size. + +```sh +.venv-tools/bin/hf download Edge0/Edge0-8B-A1B-preview \ + --local-dir Models/Edge0-8B-A1B-preview +``` + +35B is a layout change, not a change to the weights. The published checkpoint stores each expert as nine scattered reads across a 19 GB file. On an iPhone that is too slow, so the tools below rewrite each expert into one sequential read and write one file per layer. Keep the original download outside `Models/`. Only `Models/repacked/` is what the app should contain. A 35B build is about 20 GB. + +```sh +mkdir -p checkpoints +.venv-tools/bin/hf download Edge0/Edge0-35B-A3B-preview \ + --local-dir checkpoints/Edge0-35B-A3B-preview + +SOURCE=checkpoints/Edge0-35B-A3B-preview +OUTPUT=Models/repacked + +.venv-tools/bin/python tools/repack_experts.py pack "$SOURCE" "$OUTPUT" +.venv-tools/bin/python tools/repack_experts.py resident "$SOURCE" "$OUTPUT" +.venv-tools/bin/python tools/repack_experts.py verify "$SOURCE" "$OUTPUT" +.venv-tools/bin/python tools/convert_tokenizer.py "$SOURCE" "$OUTPUT/tokenizer.bin" +.venv-tools/bin/python tools/convert_pregate.py \ + "$SOURCE/prerouter_edge0_35b.safetensors" \ + "$OUTPUT/prerouter-stacked.safetensors" +cp "$SOURCE/lora_edge0_35b.safetensors" "$OUTPUT/" +``` + +`Models/repacked/` must contain `resident.safetensors`, `tokenizer.bin`, and `experts-L00.bin` through `experts-L39.bin`. You can pack either model, or both. + +Checkpoints: + +- [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) +- [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) + +### Build and install + +`com.example.edge0phone` in the project is a placeholder bundle identifier. Before building or signing, replace it with a unique bundle identifier that you control. In Xcode, open `Edge0PhoneProbe.xcodeproj`, select the **Edge0Phone** target, then change **Signing & Capabilities → Bundle Identifier**. Choose your Apple Developer Team, connect the iPhone, and Run. + +To set the team from a file instead, copy `Config/Local.xcconfig.example` to `Config/Local.xcconfig` and fill in `DEVELOPMENT_TEAM`. + +The same build from the command line: + +```sh +DEVICE_ID= +TEAM_ID= +BUNDLE_ID=com.yourname.edge0phone # replace with your own unique bundle identifier + +xcodebuild -project Edge0PhoneProbe.xcodeproj \ + -scheme Edge0Phone \ + -destination "id=$DEVICE_ID" \ + -derivedDataPath .build-device \ + -allowProvisioningUpdates \ + -skipPackagePluginValidation \ + CODE_SIGN_STYLE=Automatic \ + DEVELOPMENT_TEAM="$TEAM_ID" \ + PRODUCT_BUNDLE_IDENTIFIER="$BUNDLE_ID" \ + build + +xcrun devicectl device install app \ + --device "$DEVICE_ID" \ + .build-device/Build/Products/Debug-iphoneos/Edge0Phone.app +``` + +The weights are inside that `.app`. There is no separate copy onto the phone. + +### Test + +`scripts/test.sh` builds the MLX Metal library with Xcode, then runs the unit tests. The first run resolves the MLX Swift package. + +```sh +scripts/test.sh +EDGE0_TEST_GPU=1 scripts/test.sh +``` + +### Use the app + +Launch the app and choose **Edge0 8B** or **Edge0 35B**. A model that was not in `Models/` when you built shows **Not included** and cannot be selected. Type a question, or tap one of the suggestions. The line under a reply shows the token count, time to first token, prefill, decode, and peak memory. Switching models unloads the one that is loaded. New chat clears the conversation. + +## Performance + +| Model | Device | OS | Prefill | TTFT | Decode | +| --- | --- | --- | --- | --- | --- | +| 8B | iPhone 16 Pro | iOS 26.6.2 | 7.2 tok/s | 3.6 s | 10.9 tok/s | +| 35B | iPhone 16 Pro | iOS 26.6.2 | 4.9 tok/s | 2.1 s | 6.4 tok/s | + +## Technical Details + +The app is SwiftUI. `Edge0MLX` runs the model. `Edge0Core` reads the checkpoint, tokenizer, and routing index. Math goes through MLX Swift on Metal. The deployment target is iOS 17. + +Edge0 8B is 24 layers, alternating Kimi Delta Attention and multi-head latent attention, Top-8, INT4. Prompt prefill runs a whole layer at a time, in chunks of up to 2048 tokens. The first token is greedy. Later tokens use temperature 0.7, top-k 64, top-p 0.95, and a repetition penalty of 1.1. + +Edge0 35B is 40 layers, 256 experts per layer, Top-4. Resident weights and the per-layer expert files are read separately. The prefix cache is stored in Application Support, because the app bundle is read-only. The two models are never loaded at the same time. diff --git a/ios/Sources/Edge0Core/Edge0Configuration8B.swift b/ios/Sources/Edge0Core/Edge0Configuration8B.swift new file mode 100644 index 0000000..2047695 --- /dev/null +++ b/ios/Sources/Edge0Core/Edge0Configuration8B.swift @@ -0,0 +1,99 @@ +import Foundation + +/// Minimal configuration surface needed for the Edge0 8B port. +/// Mirrors Edge0/Edge0-8B-A1B-preview config.json. +public struct Edge0Configuration8B: Codable, Sendable, Equatable { + public let hiddenSize: Int + public let hiddenLayers: Int + public let intermediateSize: Int + public let attentionHeads: Int + public let keyValueHeads: Int + public let headDim: Int + public let vocabularySize: Int + public let maxPositionEmbeddings: Int + public let rmsNormEps: Float + + public let layerGroupSize: Int + public let firstKDenseReplace: Int + public let shortConvKernelSize: Int + public let kdaSafeGate: Bool + public let kdaLowerBound: Float + + public let qLoraRank: Int + public let kvLoraRank: Int + public let qkNopeHeadDim: Int + public let qkRopeHeadDim: Int + public let vHeadDim: Int + public let ropeTheta: Float + public let ropeInterleave: Bool + + public let numExperts: Int + public let expertsPerToken: Int + public let numSharedExperts: Int + public let moeIntermediateSize: Int + public let sharedExpertIntermediateSize: Int + public let nGroup: Int + public let topkGroup: Int + public let normTopkProb: Bool + public let routedScalingFactor: Float + public let routerHasExpertBias: Bool + + public let quantization: Quantization + + public struct Quantization: Codable, Sendable, Equatable { + public let groupSize: Int + public let bits: Int + public let mode: String + + enum CodingKeys: String, CodingKey { + case groupSize = "group_size" + case bits + case mode + } + } + + enum CodingKeys: String, CodingKey { + case hiddenSize = "hidden_size" + case hiddenLayers = "num_hidden_layers" + case intermediateSize = "intermediate_size" + case attentionHeads = "num_attention_heads" + case keyValueHeads = "num_key_value_heads" + case headDim = "head_dim" + case vocabularySize = "vocab_size" + case maxPositionEmbeddings = "max_position_embeddings" + case rmsNormEps = "rms_norm_eps" + case layerGroupSize = "layer_group_size" + case firstKDenseReplace = "first_k_dense_replace" + case shortConvKernelSize = "short_conv_kernel_size" + case kdaSafeGate = "kda_safe_gate" + case kdaLowerBound = "kda_lower_bound" + case qLoraRank = "q_lora_rank" + case kvLoraRank = "kv_lora_rank" + case qkNopeHeadDim = "qk_nope_head_dim" + case qkRopeHeadDim = "qk_rope_head_dim" + case vHeadDim = "v_head_dim" + case ropeTheta = "rope_theta" + case ropeInterleave = "rope_interleave" + case numExperts = "num_experts" + case expertsPerToken = "num_experts_per_tok" + case numSharedExperts = "num_shared_experts" + case moeIntermediateSize = "moe_intermediate_size" + case sharedExpertIntermediateSize = "moe_shared_expert_intermediate_size" + case nGroup = "n_group" + case topkGroup = "topk_group" + case normTopkProb = "norm_topk_prob" + case routedScalingFactor = "routed_scaling_factor" + case routerHasExpertBias = "moe_router_enable_expert_bias" + case quantization + } + + public static func load(from url: URL) throws -> Self { + let data = try Data(contentsOf: url) + return try JSONDecoder().decode(Self.self, from: data) + } + + public func isMLALayer(_ index: Int) -> Bool { + let full = hiddenLayers / layerGroupSize * layerGroupSize + return (index + 1) % layerGroupSize == 0 || index >= full + } +} diff --git a/ios/Sources/Edge0Core/Edge0Tokenizer8B.swift b/ios/Sources/Edge0Core/Edge0Tokenizer8B.swift new file mode 100644 index 0000000..dd1a9eb --- /dev/null +++ b/ios/Sources/Edge0Core/Edge0Tokenizer8B.swift @@ -0,0 +1,186 @@ +import Foundation + +public enum Edge0TokenizerError8B: Error, LocalizedError { + case invalid(String) + + public var errorDescription: String? { + switch self { + case .invalid(let message): message + } + } +} + +/// Native Swift implementation of the checkpoint's NFC + regex + ByteLevel BPE pipeline. +/// It reads `tokenizer.json` from the model folder to use the checkpoint vocabulary and +/// merge table directly. +public final class Edge0Tokenizer8B: @unchecked Sendable { + private struct FileFormat: Decodable { + struct AddedToken: Decodable { + let id: Int + let content: String + } + + struct Model: Decodable { + let vocab: [String: Int] + let merges: [[String]] + } + + let added_tokens: [AddedToken] + let model: Model + } + + private struct Pair: Hashable { + let left: String + let right: String + } + + private static let splitPattern = #"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"# + + private let vocabulary: [String: Int] + private let tokensByID: [Int: String] + private let mergeRanks: [Pair: Int] + private let addedTokens: [(content: String, id: Int)] + private let byteEncoder: [UInt8: String] + private let byteDecoder: [Character: UInt8] + private let splitRegex: NSRegularExpression + + public let endOfTurnTokenID: Int + + public init(contentsOf url: URL) throws { + let payload = try JSONDecoder().decode(FileFormat.self, from: Data(contentsOf: url)) + vocabulary = payload.model.vocab + tokensByID = Dictionary(uniqueKeysWithValues: payload.model.vocab.map { ($0.value, $0.key) }) + mergeRanks = Dictionary(uniqueKeysWithValues: payload.model.merges.enumerated().compactMap { rank, merge in + guard merge.count == 2 else { return nil } + return (Pair(left: merge[0], right: merge[1]), rank) + }) + addedTokens = payload.added_tokens + .map { ($0.content, $0.id) } + .sorted { lhs, rhs in lhs.content.count > rhs.content.count } + guard let eos = payload.added_tokens.first(where: { $0.content == "<|role_end|>" })?.id else { + throw Edge0TokenizerError8B.invalid("tokenizer.json is missing <|role_end|>") + } + endOfTurnTokenID = eos + + let maps = Self.makeByteMaps() + byteEncoder = maps.encoder + byteDecoder = maps.decoder + splitRegex = try NSRegularExpression(pattern: Self.splitPattern) + } + + public func encode(_ text: String) throws -> [Int] { + let normalized = text.precomposedStringWithCanonicalMapping + var ids: [Int] = [] + var plain = "" + var cursor = normalized.startIndex + + func flushPlain() throws { + guard !plain.isEmpty else { return } + ids.append(contentsOf: try encodePlain(plain)) + plain.removeAll(keepingCapacity: true) + } + + while cursor < normalized.endIndex { + let suffix = normalized[cursor...] + if let token = addedTokens.first(where: { suffix.hasPrefix($0.content) }) { + try flushPlain() + ids.append(token.id) + cursor = normalized.index(cursor, offsetBy: token.content.count) + } else { + let next = normalized.index(after: cursor) + plain.append(contentsOf: normalized[cursor.. String { + var bytes: [UInt8] = [] + for id in ids { + guard let token = tokensByID[id] ?? addedTokens.first(where: { $0.id == id })?.content else { + throw Edge0TokenizerError8B.invalid("Unknown token ID \(id)") + } + if skipSpecialTokens, addedTokens.contains(where: { $0.id == id }) { continue } + for character in token { + if let byte = byteDecoder[character] { + bytes.append(byte) + } else { + bytes.append(contentsOf: String(character).utf8) + } + } + } + return String(decoding: bytes, as: UTF8.self) + } + + private func encodePlain(_ text: String) throws -> [Int] { + let range = NSRange(text.startIndex.. [String] { + var pieces = token.map(String.init) + while pieces.count > 1 { + var selected: Pair? + var selectedRank = Int.max + for index in 0..<(pieces.count - 1) { + let pair = Pair(left: pieces[index], right: pieces[index + 1]) + if let rank = mergeRanks[pair], rank < selectedRank { + selected = pair + selectedRank = rank + } + } + guard let selected else { break } + + var merged: [String] = [] + var index = 0 + while index < pieces.count { + if index + 1 < pieces.count, + pieces[index] == selected.left, + pieces[index + 1] == selected.right { + merged.append(selected.left + selected.right) + index += 2 + } else { + merged.append(pieces[index]) + index += 1 + } + } + pieces = merged + } + return pieces + } + + private static func makeByteMaps() -> (encoder: [UInt8: String], decoder: [Character: UInt8]) { + var bytes = Array(UInt8(33)...UInt8(126)) + bytes += Array(UInt8(161)...UInt8(172)) + bytes += Array(UInt8(174)...UInt8(255)) + var scalars = bytes.map(Int.init) + var extra = 0 + for value in 0...255 where !bytes.contains(UInt8(value)) { + bytes.append(UInt8(value)) + scalars.append(256 + extra) + extra += 1 + } + var encoder: [UInt8: String] = [:] + var decoder: [Character: UInt8] = [:] + for (byte, scalarValue) in zip(bytes, scalars) { + let character = Character(String(UnicodeScalar(scalarValue)!)) + encoder[byte] = String(character) + decoder[character] = byte + } + return (encoder, decoder) + } +} diff --git a/ios/Sources/Edge0Core/ExpertTensorStore.swift b/ios/Sources/Edge0Core/ExpertTensorStore.swift new file mode 100644 index 0000000..ae7412f --- /dev/null +++ b/ios/Sources/Edge0Core/ExpertTensorStore.swift @@ -0,0 +1,60 @@ +import Foundation + +public enum ExpertProjection: String, CaseIterable, Sendable { + case up = "up_proj" + case gate = "gate_proj" + case down = "down_proj" +} + +public enum QuantPart: String, CaseIterable, Sendable { + case weight + case scales + case biases +} + +public struct ExpertTensorKey: Hashable, Sendable { + public let layer: Int + public let expert: Int + public let projection: ExpertProjection + public let part: QuantPart +} + +public struct ExpertTensorSlice: @unchecked Sendable { + public let descriptor: TensorDescriptor + public let pointer: UnsafeMutableRawPointer +} + +/// Index + mmap layer for the Edge0 checkpoint. No MLX dependency here. +/// The MLX adapter can turn these slices into arrays only when an expert is selected. +public final class ExpertTensorStore: @unchecked Sendable { + public let index: SafetensorsIndex + public let mappedFile: MappedFile + public let expertCount: Int + + public init(modelURL: URL, expertCount: Int = 128) throws { + self.index = try SafetensorsIndex(fileURL: modelURL) + self.mappedFile = try MappedFile(url: modelURL) + self.expertCount = expertCount + } + + public func stackedTensorName(layer: Int, projection: ExpertProjection, part: QuantPart) -> String { + "model.layers.\(layer).mlp.experts.\(projection.rawValue).\(part.rawValue)" + } + + public func slice( + layer: Int, + expert: Int, + projection: ExpertProjection, + part: QuantPart + ) throws -> ExpertTensorSlice { + precondition((0.. Selection { + precondition(nGroup > 0) + precondition(logits.count / nGroup >= 2) + precondition(topK <= topkGroup * (logits.count / nGroup)) + precondition(logits.count == expertBias.count) + precondition(logits.count % nGroup == 0) + precondition(topK > 0 && topK <= logits.count) + precondition(topkGroup > 0 && topkGroup <= nGroup) + + let raw = logits.map { 1 / (1 + Foundation.exp(-$0)) } + var selection = zip(raw, expertBias).map(+) + let perGroup = logits.count / nGroup + + let groupScores: [(group: Int, score: Float)] = (0..).prefix(2) + return (group, bestTwo.reduce(0, +)) + } + let surviving = Set(groupScores.sorted { $0.score > $1.score }.prefix(topkGroup).map(\.group)) + + for group in 0.. selection[$1] + }.prefix(topK).map { $0 } + + var weights = indices.map { raw[$0] } + if normalize { + let denominator = weights.reduce(0, +) + 1e-20 + weights = weights.map { $0 / denominator } + } + weights = weights.map { $0 * routedScale } + return Selection(indices: indices, weights: weights) + } +} diff --git a/ios/Sources/Edge0Core/MappedFile.swift b/ios/Sources/Edge0Core/MappedFile.swift new file mode 100644 index 0000000..51d1853 --- /dev/null +++ b/ios/Sources/Edge0Core/MappedFile.swift @@ -0,0 +1,123 @@ +import Foundation +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#endif + +public enum MappedFileError: Error, LocalizedError { + case openFailed(String) + case statFailed(String) + case mapFailed(String) + case outOfBounds(Range) + + public var errorDescription: String? { + switch self { + case .openFailed(let p): return "Could not open \(p)." + case .statFailed(let p): return "Could not stat \(p)." + case .mapFailed(let p): return "Could not mmap \(p)." + case .outOfBounds(let r): return "Mapped byte range out of bounds: \(r)." + } + } +} + +/// Read-only virtual-memory mapping of a model file. +/// mmap does not make the whole file resident; pages are faulted in on demand. +public final class MappedFile: @unchecked Sendable { + public let url: URL + public let count: Int + private let fd: Int32 + private let base: UnsafeMutableRawPointer + + public init(url: URL) throws { + self.url = url + let path = url.path + let descriptor = open(path, O_RDONLY) + guard descriptor >= 0 else { throw MappedFileError.openFailed(path) } + + var st = stat() + guard fstat(descriptor, &st) == 0 else { + close(descriptor) + throw MappedFileError.statFailed(path) + } + let size = Int(st.st_size) + guard size > 0 else { + close(descriptor) + throw MappedFileError.statFailed(path) + } + + let mapping = mmap(nil, size, PROT_READ, MAP_PRIVATE, descriptor, 0) + guard mapping != MAP_FAILED, let mapping else { + close(descriptor) + throw MappedFileError.mapFailed(path) + } + + self.fd = descriptor + self.count = size + self.base = mapping + } + + deinit { + munmap(base, count) + close(fd) + } + + public func pointer(to range: Range) throws -> UnsafeMutableRawPointer { + guard range.lowerBound >= 0, range.upperBound <= count else { + throw MappedFileError.outOfBounds(range) + } + return base.advanced(by: range.lowerBound) + } + + public func bytes(in range: Range) throws -> UnsafeRawBufferPointer { + let ptr = try pointer(to: range) + return UnsafeRawBufferPointer(start: UnsafeRawPointer(ptr), count: range.count) + } + + /// Advisory only; ignored if the OS declines it. + public func adviseSequential() { + #if canImport(Darwin) + _ = madvise(base, count, MADV_SEQUENTIAL) + #endif + } + + /// Official startup prewarm hint. Advisory only; iOS remains free to + /// limit page-cache residency under memory pressure. + public func adviseWillNeed() { + #if canImport(Darwin) + _ = madvise(base, count, MADV_WILLNEED) + #endif + } + + /// Advisory readahead for a checkpoint subrange. Streaming experts use + /// this before copying their tensor slices so page faults can be serviced + /// in bulk rather than serially on the decode critical path. + public func adviseWillNeed(range: Range) { + #if canImport(Darwin) + guard range.lowerBound >= 0, range.upperBound <= count, + !range.isEmpty else { return } + let pageSize = max(4096, Int(sysconf(_SC_PAGESIZE))) + let start = range.lowerBound / pageSize * pageSize + let end = min( + count, + ((range.upperBound + pageSize - 1) / pageSize) * pageSize) + _ = madvise(base.advanced(by: start), end - start, MADV_WILLNEED) + #endif + } + + /// Official `SafetensorsMmap.seq_read`: force a real sequential pass over + /// the checkpoint. `madvise` alone does not reliably fault every page in. + public func sequentialRead(chunkSize: Int = 1 << 24) { + let size = max(4096, chunkSize) + let buffer = UnsafeMutableRawPointer.allocate( + byteCount: size, alignment: MemoryLayout.alignment) + defer { buffer.deallocate() } + var offset = 0 + while offset < count { + let requested = min(size, count - offset) + let bytesRead = pread(fd, buffer, requested, off_t(offset)) + guard bytesRead > 0 else { break } + offset += bytesRead + } + } +} diff --git a/ios/Sources/Edge0Core/SafetensorsIndex.swift b/ios/Sources/Edge0Core/SafetensorsIndex.swift new file mode 100644 index 0000000..be08109 --- /dev/null +++ b/ios/Sources/Edge0Core/SafetensorsIndex.swift @@ -0,0 +1,152 @@ +import Foundation + +public enum SafetensorsError: Error, LocalizedError { + case fileTooSmall + case invalidHeaderLength(UInt64) + case invalidHeader + case malformedTensor(String) + case unsupportedDType(String) + case tensorNotFound(String) + case invalidExpertAxis(String) + + public var errorDescription: String? { + switch self { + case .fileTooSmall: return "Safetensors file is too small." + case .invalidHeaderLength(let n): return "Invalid safetensors header length: \(n)." + case .invalidHeader: return "Safetensors JSON header is invalid." + case .malformedTensor(let name): return "Malformed tensor entry: \(name)." + case .unsupportedDType(let dtype): return "Unsupported safetensors dtype: \(dtype)." + case .tensorNotFound(let name): return "Tensor not found: \(name)." + case .invalidExpertAxis(let name): return "Tensor has no valid expert axis: \(name)." + } + } +} + +public enum SafetensorsDType: String, Sendable, Codable { + case bool = "BOOL" + case u8 = "U8" + case i8 = "I8" + case i16 = "I16" + case u16 = "U16" + case i32 = "I32" + case u32 = "U32" + case i64 = "I64" + case u64 = "U64" + case f16 = "F16" + case bf16 = "BF16" + case f32 = "F32" + case f64 = "F64" + + public var byteWidth: Int { + switch self { + case .bool, .u8, .i8: 1 + case .i16, .u16, .f16, .bf16: 2 + case .i32, .u32, .f32: 4 + case .i64, .u64, .f64: 8 + } + } +} + +public struct TensorDescriptor: Sendable, Equatable { + public let name: String + public let dtype: SafetensorsDType + public let shape: [Int] + /// Absolute byte range within the safetensors file. + public let byteRange: Range + + public var byteCount: Int { byteRange.count } + + /// Returns a contiguous slice along axis 0. Edge0 stores routed experts stacked on axis 0. + public func axisZeroSlice(_ index: Int) throws -> TensorDescriptor { + guard let count = shape.first, count > 0, index >= 0, index < count else { + throw SafetensorsError.invalidExpertAxis(name) + } + guard byteCount % count == 0 else { + throw SafetensorsError.invalidExpertAxis(name) + } + let stride = byteCount / count + let start = byteRange.lowerBound + index * stride + return TensorDescriptor( + name: "\(name)[\(index)]", + dtype: dtype, + shape: Array(shape.dropFirst()), + byteRange: start..<(start + stride) + ) + } +} + +public struct SafetensorsIndex: Sendable { + public let fileURL: URL + public let headerLength: Int + public let dataOffset: Int + public let tensors: [String: TensorDescriptor] + + public init(fileURL: URL) throws { + self.fileURL = fileURL + let handle = try FileHandle(forReadingFrom: fileURL) + defer { try? handle.close() } + + let prefix = try handle.read(upToCount: 8) ?? Data() + guard prefix.count == 8 else { throw SafetensorsError.fileTooSmall } + let headerLength64 = prefix.withUnsafeBytes { raw -> UInt64 in + raw.loadUnaligned(as: UInt64.self).littleEndian + } + guard headerLength64 <= UInt64(Int.max) else { + throw SafetensorsError.invalidHeaderLength(headerLength64) + } + let headerLength = Int(headerLength64) + guard headerLength > 0 else { throw SafetensorsError.invalidHeaderLength(headerLength64) } + + let headerData = try handle.read(upToCount: headerLength) ?? Data() + guard headerData.count == headerLength else { throw SafetensorsError.fileTooSmall } + guard let json = try JSONSerialization.jsonObject(with: headerData) as? [String: Any] else { + throw SafetensorsError.invalidHeader + } + + let base = 8 + headerLength + var parsed: [String: TensorDescriptor] = [:] + parsed.reserveCapacity(json.count) + + for (name, value) in json where name != "__metadata__" { + guard + let entry = value as? [String: Any], + let dtypeString = entry["dtype"] as? String, + let dtype = SafetensorsDType(rawValue: dtypeString), + let shapeNumbers = entry["shape"] as? [NSNumber], + let offsets = entry["data_offsets"] as? [NSNumber], + offsets.count == 2 + else { + if let entry = value as? [String: Any], let dtypeString = entry["dtype"] as? String, + SafetensorsDType(rawValue: dtypeString) == nil { + throw SafetensorsError.unsupportedDType(dtypeString) + } + throw SafetensorsError.malformedTensor(name) + } + + let shape = shapeNumbers.map(\.intValue) + let relativeStart = offsets[0].intValue + let relativeEnd = offsets[1].intValue + guard relativeStart >= 0, relativeEnd >= relativeStart else { + throw SafetensorsError.malformedTensor(name) + } + + parsed[name] = TensorDescriptor( + name: name, + dtype: dtype, + shape: shape, + byteRange: (base + relativeStart)..<(base + relativeEnd) + ) + } + + self.headerLength = headerLength + self.dataOffset = base + self.tensors = parsed + } + + public subscript(_ name: String) -> TensorDescriptor? { tensors[name] } + + public func tensor(named name: String) throws -> TensorDescriptor { + guard let tensor = tensors[name] else { throw SafetensorsError.tensorNotFound(name) } + return tensor + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift b/ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift new file mode 100644 index 0000000..7cd6966 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift @@ -0,0 +1,122 @@ +import Foundation + +/// The text-only slice of this checkpoint's chat template. +/// +/// The shipped `chat_template.jinja` is 7,764 characters of Jinja with macros, image and +/// video handling, tool calls and vision counters. None of that is reachable from a +/// text-only app, and running a Jinja interpreter on device to reach a fixed string would +/// be a large amount of machinery for no behaviour. +/// +/// What the text path actually renders is regular: +/// +/// ``` +/// <|im_start|>{role}\n{content}<|im_end|>\n per message +/// <|im_start|>assistant\n\n generation prompt (thinking on) +/// <|im_start|>assistant\n generation prompt (thinking off) +/// ``` +/// +/// **Scope is the point, not a shortcut.** Vision and tool-call messages are out — and +/// they fail loudly rather than rendering something that looks plausible, because a +/// silently dropped image would produce a coherent answer about nothing. +/// +/// The trailing `\n` is the checkpoint's default: replies open inside a reasoning +/// block. Passing `thinking: false` omits it so the model writes the answer directly — +/// fewer tokens before anything useful appears, at the cost of losing that reasoning pass. +enum ChatTemplate { + + struct Message { + let role: String + let content: String + } + + enum Failure: LocalizedError { + case unsupportedRole(String) + + var errorDescription: String? { + switch self { + case .unsupportedRole(let role): + return "chat template supports system/user/assistant only, not '\(role)'" + } + } + } + + static let supportedRoles: Set = ["system", "user", "assistant"] + static let userOpening = "<|im_start|>user\n" + + /// The opening tag every assistant turn is generated under. + static func assistantOpening(thinking: Bool) -> String { + thinking + ? "<|im_start|>assistant\n\n" + : "<|im_start|>assistant\n" + } + + static func render( + _ messages: [Message], + addGenerationPrompt: Bool = true, + thinking: Bool = true + ) throws -> String { + var out = "" + for message in messages { + guard supportedRoles.contains(message.role) else { + throw Failure.unsupportedRole(message.role) + } + // **Past assistant turns are reopened with the same thinking mode.** + // Generation runs under `assistantOpening`, so a thinking-on reply begins + // inside a reasoning block and contains its ``. Rendering that text + // back without the opening tag put a dangling close tag in the history — + // well-formed-looking output describing a prompt the model never saw. + // + // It also makes replay exact, which is what lets a conversation's state be + // extended instead of rebuilt. + if message.role == "assistant" { + out += assistantOpening(thinking: thinking) + "\(message.content)<|im_end|>\n" + } else { + out += message.role == "user" + ? userOpening + "\(message.content)<|im_end|>\n" + : "<|im_start|>\(message.role)\n\(message.content)<|im_end|>\n" + } + } + if addGenerationPrompt { + out += assistantOpening(thinking: thinking) + } + return out + } + + /// What must be fed to continue a conversation whose state already ends inside the + /// previous assistant turn. + /// + /// This is the append path: it closes the open assistant turn, adds the new question, + /// and reopens for the answer. Nothing before it is re-rendered, so nothing before it + /// is recomputed. + static func continuation(user: String, thinking: Bool = true) -> String { + "<|im_end|>\n" + userOpening + "\(user)<|im_end|>\n" + assistantOpening(thinking: thinking) + } +} + +/// When to stop generating. +/// +/// Two conditions, and both are needed: `eos_token_id` is a **list** here — 248046 +/// (`<|im_end|>`) and 248044 (`<|endoftext|>`) — so checking only the first lets the +/// model run past the end of its turn. The length bound is the backstop for the case +/// where neither is ever emitted, which is what a mis-specified prompt produces. +struct StopCondition { + let endTokens: Set + let maximumTokens: Int + + init(endTokens: [Int32] = [248046, 248044], maximumTokens: Int = 256) { + self.endTokens = Set(endTokens) + self.maximumTokens = maximumTokens + } + + enum Outcome: Equatable { + case keepGoing + case hitEndToken(Int32) + case hitLimit + } + + func evaluate(token: Int32, produced: Int) -> Outcome { + if endTokens.contains(token) { return .hitEndToken(token) } + if produced >= maximumTokens { return .hitLimit } + return .keepGoing + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift b/ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift new file mode 100644 index 0000000..084e5c9 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift @@ -0,0 +1,60 @@ +import CoreML +import Foundation +import MLX + +/// Optional Core ML vocabulary projection used by the 35B path. +/// +/// The model is discovered beside the streamed checkpoint rather than bundled with the +/// app: it is a derived 273 MiB asset and the existing MLX path remains the fallback +/// when it is absent or fails parity. Argmax stays in the graph, so only one Int32 token +/// crosses back from Core ML. +final class CoreMLVocabularyHead: @unchecked Sendable { + static let directoryName = "lm-head.mlmodelc" + static let width = 2_048 + + private let model: MLModel + + init(directory: URL) throws { + let url = directory.appendingPathComponent(Self.directoryName, isDirectory: true) + let configuration = MLModelConfiguration() + // Keep Core ML off the GPU so its projection does not compete with the MLX/Metal + // decoder. Unsupported operations can fall back to the CPU. + configuration.computeUnits = .cpuAndNeuralEngine + model = try MLModel(contentsOf: url, configuration: configuration) + } + + func predict(_ hidden: MLXArray) async throws -> Int32 { + let values = hidden.asType(.float16).asArray(Float16.self) + guard values.count == Self.width else { + throw Failure.badHiddenSize(values.count) + } + + let input = try MLMultiArray( + shape: [1, NSNumber(value: Self.width)], dataType: .float16) + values.withUnsafeBytes { source in + input.dataPointer.copyMemory(from: source.baseAddress!, byteCount: source.count) + } + let provider = try MLDictionaryFeatureProvider(dictionary: [ + "hidden_state": MLFeatureValue(multiArray: input) + ]) + let output = try await model.prediction(from: provider) + guard let token = output.featureValue(for: "token")?.multiArrayValue, + token.count == 1 + else { throw Failure.missingToken } + return token[0].int32Value + } + + private enum Failure: LocalizedError { + case badHiddenSize(Int) + case missingToken + + var errorDescription: String? { + switch self { + case .badHiddenSize(let count): + "Core ML lm_head expected \(CoreMLVocabularyHead.width) values, got \(count)" + case .missingToken: + "Core ML lm_head did not return token" + } + } + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift b/ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift new file mode 100644 index 0000000..518c58c --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift @@ -0,0 +1,284 @@ +import Darwin +import Foundation +import MLX + +/// Product-facing wrapper around the streaming Edge0 35B implementation. +/// +/// This engine deliberately lives beside `Edge0ChatEngine`: selecting 35B creates only +/// this model, while selecting 8B continues through the existing engine unchanged. +public final class Edge0ChatEngine35B: @unchecked Sendable { + public static let modelFolderName = "repacked" + + private let directory: URL + private let model: Edge0Model35B + private let tokenizer: Tokenizer + private let instructions: String + private let systemPrefix: [Int32] + private let state = Edge0Model35B.State() + + private var processed: [Int32] = [] + private var cachedPrefix: PrefixCache? + + public init( + modelURL: URL, + instructions: String = "You are a helpful assistant.", + progress: (String) -> Void = { _ in } + ) throws { + directory = modelURL + self.instructions = instructions + + // Configure MLX memory use for the selected 35B runtime. + let highMem = CommandLine.arguments.contains("--high-mem-35b") + Memory.cacheLimit = (highMem ? 1536 : 384) * 1_048_576 + print(String( + format: "[mem] cacheLimit=%dMiB highMem=%@ available=%dMiB", + Memory.cacheLimit / 1_048_576, + highMem ? "true" : "false", + os_proc_available_memory() / 1_048_576)) + progress("mapping Edge0 35B weights") + var configuration = Edge0Model35B.Configuration() + configuration.round6 = CommandLine.arguments.contains("--round6") + if configuration.round6 { + configuration.activeExperts = 2 + } + configuration.retainAllResidentWeights = highMem + model = try Edge0Model35B(directory: modelURL, configuration: configuration) + guard model.availableLayers == model.configuration.layers else { + throw Edge035BFailure.incompleteExperts( + found: model.availableLayers, expected: model.configuration.layers) + } + progress("loading Edge0 35B tokenizer") + tokenizer = try Tokenizer(url: modelURL.appendingPathComponent("tokenizer.bin")) + + if instructions.isEmpty { + systemPrefix = [] + } else { + let text = try ChatTemplate.render( + [.init(role: "system", content: instructions)], + addGenerationPrompt: false) + // The first user role marker is fixed too. Persisting its state saves those + // tokens on every new conversation; `PrefixCache.matches` still verifies the + // encoded prefix against the complete prompt before it is trusted. + systemPrefix = tokenizer.encode(text + ChatTemplate.userOpening) + } + cachedPrefix = PrefixCache.read(from: Self.prefixURL( + directory: modelURL, fingerprint: model.configuration.fingerprint)) + try model.prepareForFirstTurn() + progress("Edge0 35B ready") + } + + public func reset() { + state.reset() + processed.removeAll(keepingCapacity: true) + } + + public func reply( + to userText: String, + maxTokens: Int = 512, + thinking: Bool = true, + onText: @escaping @Sendable (String) -> Void = { _ in }, + shouldContinue: @escaping @Sendable () -> Bool = { true } + ) async throws -> Edge0GenerationResult { + let prompt = userText.trimmingCharacters(in: .whitespacesAndNewlines) + guard !prompt.isEmpty else { throw Edge035BFailure.invalid("Message cannot be empty") } + guard maxTokens > 0 else { throw Edge035BFailure.invalid("maxTokens must be positive") } + + let continuing = !processed.isEmpty + let ids: [Int32] + if continuing { + ids = tokenizer.encode(ChatTemplate.continuation(user: prompt, thinking: thinking)) + } else { + var messages: [ChatTemplate.Message] = [] + if !instructions.isEmpty { + messages.append(.init(role: "system", content: instructions)) + } + messages.append(.init(role: "user", content: prompt)) + ids = tokenizer.encode(try ChatTemplate.render(messages, thinking: thinking)) + } + guard !ids.isEmpty else { throw Edge035BFailure.invalid("Tokenizer returned an empty prompt") } + print("[35b] thinking=\(thinking ? "on" : "off")") + + GPU.resetPeakMemory() + let started = Date() + let thermalAtStart = ProcessInfo.processInfo.thermalState + var start = 0 + if !continuing { + state.reset() + processed.removeAll(keepingCapacity: true) + if let cache = cachedPrefix, cache.matches(ids) { + restore(cache.state, into: state) + processed = cache.tokens + start = cache.tokens.count + } + } + + var next: Int32 = 0 + // Later turns already have attention weights resident. Keep that cache on + // for a continued prefill so those projections are not copied again. + // The first prefill still copies each of them once; retaining earlier + // only adds dirty pages with nothing to reuse inside that single pass. + if continuing { model.beginDecodeCaching() } + model.profile.reset() + // An older cache may contain only the system message. Restore it, compute just + // the newly-added user opening, then atomically replace it with the longer cache. + let capturePoint = !continuing && start < systemPrefix.count + && ids.starts(with: systemPrefix) ? systemPrefix.count : nil + var index = start + while index < ids.count { + guard shouldContinue() else { return try cancelled() } + var size = min(Edge0Model35B.maximumBatch, ids.count - index) + if let capturePoint, index < capturePoint { + size = min(size, capturePoint - index) + } + let slice = Array(ids[index ..< index + size]) + let isLast = index + size == ids.count + let logits = try model.step( + tokens: MLXArray(slice, [1, size]), state: state, needsLogits: isLast) + if isLast, let logits { next = try await model.greedy(logits) } + processed.append(contentsOf: slice) + index += size + + if let capturePoint, index == capturePoint { + let captured = PrefixCache(tokens: systemPrefix, state: state.snapshot()) + cachedPrefix = captured + captured.write(to: Self.prefixURL( + directory: directory, fingerprint: model.configuration.fingerprint)) + } + } + let prefillFinished = Date() + let prefillSeconds = prefillFinished.timeIntervalSince(started) + let prefillTokens = ids.count - start + print(String( + format: "[35b prefill] tokens=%d total=%.2fs rate=%.2f tok/s %@", + prefillTokens, + prefillSeconds, + Double(prefillTokens) / max(0.000_001, prefillSeconds), + model.profile.report)) + model.beginDecodeCaching() + model.profile.reset() + + let stop = StopCondition(maximumTokens: maxTokens) + var raw = "" + var shown = "" + var produced = 0 + var firstTokenAt: Date? + var decodeStartedAt: Date? + + while true { + switch stop.evaluate(token: next, produced: produced) { + case .hitEndToken, .hitLimit: + let finished = Date() + let decodeSeconds = decodeStartedAt.map { + finished.timeIntervalSince($0) + } ?? 0 + let footprint = Self.physicalFootprintBytes() + print(String( + format: "[35b decode] tokens=%d total=%.2fs rate=%.2f tok/s thermal=%@→%@ mlxPeak=%.0fMB footprint=%.0fMB available=%dMB %@", + max(0, produced - 1), + decodeSeconds, + Double(max(0, produced - 1)) / max(0.000_001, decodeSeconds), + Self.thermalName(thermalAtStart), + Self.thermalName(ProcessInfo.processInfo.thermalState), + Double(Memory.peakMemory) / 1_048_576, + Double(footprint) / 1_048_576, + os_proc_available_memory() / 1_048_576, + model.profile.report)) + shown = thinking + ? ThinkingBlock.settled(raw) + : ThinkingBlock.settledReply(raw) + if !shown.isEmpty { onText(shown) } + return Edge0GenerationResult( + text: shown, + generatedTokenCount: produced, + elapsedSeconds: finished.timeIntervalSince(started), + prefillSeconds: prefillSeconds, + prefillTokensPerSecond: Double(prefillTokens) / + max(0.000_001, prefillSeconds), + timeToFirstTokenSeconds: firstTokenAt?.timeIntervalSince(started) + ?? finished.timeIntervalSince(started), + decodeTokensPerSecond: decodeStartedAt.map { + Double(max(0, produced - 1)) / + max(0.000_001, finished.timeIntervalSince($0)) + } ?? 0, + peakMemoryBytes: max(Memory.peakMemory, footprint)) + case .keepGoing: + break + } + + guard shouldContinue() else { return try cancelled() } + if firstTokenAt == nil { + firstTokenAt = Date() + decodeStartedAt = firstTokenAt + } + raw += tokenizer.decode([next]) + let visible = thinking + ? ThinkingBlock.visible(in: raw) + : ThinkingBlock.visibleReply(in: raw) + if let visible, visible != shown { + shown = visible + onText(shown) + } + produced += 1 + + let logits = try model.step(tokens: MLXArray([next], [1, 1]), state: state) + processed.append(next) + next = try await model.greedy(logits) + } + } + + private func cancelled() throws -> T { + reset() + throw CancellationError() + } + + private func restore(_ source: Edge0Model35B.State, into target: Edge0Model35B.State) { + let snapshot = source.snapshot() + target.convolution = snapshot.convolution + target.recurrent = snapshot.recurrent + target.keys = snapshot.keys + target.values = snapshot.values + target.offset = snapshot.offset + } + + private static func prefixURL(directory: URL, fingerprint: String) -> URL { + PrefixCache.fileURL(in: directory, fingerprint: fingerprint) + } + + private static func thermalName(_ state: ProcessInfo.ThermalState) -> String { + switch state { + case .nominal: "nominal" + case .fair: "fair" + case .serious: "serious" + case .critical: "critical" + @unknown default: "unknown" + } + } + + /// Process physical footprint (jetsam-relevant), not just MLX allocator peak. + private static func physicalFootprintBytes() -> Int { + var info = task_vm_info_data_t() + var count = mach_msg_type_number_t( + MemoryLayout.stride / MemoryLayout.stride) + let result = withUnsafeMutablePointer(to: &info) { + $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { + task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count) + } + } + guard result == KERN_SUCCESS else { return 0 } + return Int(info.phys_footprint) + } +} + +private enum Edge035BFailure: LocalizedError { + case incompleteExperts(found: Int, expected: Int) + case invalid(String) + + var errorDescription: String? { + switch self { + case .incompleteExperts(let found, let expected): + return "Edge0 35B expert files are incomplete (found \(found), expected \(expected))" + case .invalid(let message): + return message + } + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift b/ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift new file mode 100644 index 0000000..1b1669a --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift @@ -0,0 +1,1012 @@ +import Foundation +import MLX +import MLXFast + +/// Streaming Edge0 35B model implementation used by the iOS chat engine. +/// +/// This implementation is kept independent of the Python reference so parity checks can +/// compare separate code paths. It is not main-actor isolated because model steps perform +/// synchronous GPU and file I/O. The chat engine serializes access to one model instance. +final class Edge0Model35B: @unchecked Sendable { + + /// How many tokens may go through one forward pass. + /// + /// Sets the expert buffer ceiling: a batch of T tokens can in the worst case route to + /// T*K distinct experts. The loader grows staging storage only to the union actually + /// observed, so this ceiling no longer reserves the worst case up front. + /// + /// Storage is sized for the observed expert union rather than the worst-case T*K. + /// Top-4 routing makes the batch size a tradeoff between repeated base-weight reads + /// and the memory required for the active graph and expert union. + /// Tokens per prefill forward. Decode ignores this. + /// + /// A full-layer read is about 8 s of flash for 40 layers, independent of + /// how many tokens share it. 512 is the first size at which that fixed + /// cost can land near 60 tok/s before GPU time; the Gated DeltaNet scan + /// is what will jetsam first if 512 is too large. + static let maximumBatch = 512 + /// Below this, prefill keeps the decode-style expert union. A full-layer + /// read is 432 MiB × 40 and only pays for itself on a long chunk. + static let fullLayerPrefillMinimum = 64 + + struct Configuration { + var layers = 40 + var hidden = 2048 + var moeIntermediate = 512 + var expertsPerLayer = 256 + // The public Edge0-35B-A3B-preview adapter and prerouter were trained with K=4. + // This must stay aligned with the shipped artifacts: changing K changes both + // the MoE output and the routing history consumed by the next-step heads. + var activeExperts = 4 + /// Round-6 pregate + LoRA, trained at K=2. The two files are a pair. + var round6 = false + var groupSize = 64 + var bits = 4 + var fullAttentionInterval = 4 + var vocabulary = 248_320 + + /// When true, keep every non-expert resident tensor after first materialisation. + /// The default retains only selected tensors to limit active memory use. + var retainAllResidentWeights = false + + /// Source used to prefetch expert reads for the next decode step. + var routingGuess: RoutingGuess = .trainedHeads + + /// Number of recent routing steps used by `.recentRouting`. + var routingLookback = 1 + + enum RoutingGuess { + /// Do not prefetch expert reads. + case none + /// Prefetch the experts selected by each layer on the previous step. + case recentRouting + /// The trained pregate heads, when their file is present. Predicts the + /// *next* layer a *token* ahead, so it is the only source that can cover a + /// layer before that layer has run — and it replaces the router to do it. + case trainedHeads + } + + /// Identifies this configuration in caches whose contents depend on it. + var fingerprint: String { "full" } + + /// `(index + 1) % interval == 0` are full attention; the rest are linear. + func isLinear(_ layer: Int) -> Bool { + (layer + 1) % fullAttentionInterval != 0 + } + } + + enum Failure: LocalizedError { + case weightsMissing(URL) + case tensorMissing(String) + case layerFailed(Int) + + var errorDescription: String? { + switch self { + case .weightsMissing(let url): + return "no model at \(url.lastPathComponent)" + case .tensorMissing(let name): + return "missing tensor \(name)" + case .layerFailed(let layer): + return "layer \(layer) could not run" + } + } + } + + /// State carried between decode steps. + /// + /// Two kinds, and **they are not interchangeable**. Linear layers hold a fixed-size + /// `[1, heads, dim, dim]` float32 matrix updated in place: it does not grow with + /// context and cannot be pruned. Full-attention layers hold a KV cache that does + /// grow. Any context policy that drops old turns shortens the second and leaves the + /// first untouched — a model that appears to forget while still carrying everything. + final class State { + var convolution: [Int: MLXArray] = [:] + var recurrent: [Int: MLXArray] = [:] + var keys: [Int: MLXArray] = [:] + var values: [Int: MLXArray] = [:] + var offset = 0 + + /// Routing carried between decode steps. Per-conversation, not per-model: two + /// conversations share one set of heads but must never share the routing + /// history that feeds them. + let pregate = PregateHeads.State() + + /// Layer → the experts it ran, most recent step first, capped at the lookback. + /// Previous-step router output, used by `.recentRouting`. + var recentRouting: [Int: [[Int]]] = [:] + + /// Layer → what was asked for on its behalf. Exists so the hit rate is a + /// measurement of what was *issued* rather than a restatement of the guess: + /// `recentRouting` is overwritten as each layer runs, so comparing against it + /// would compare a layer's routing with its own. + var prefetched: [Int: [Int]] = [:] + + var byteCount: Int { + [convolution, recurrent, keys, values].reduce(0) { total, table in + total + table.values.reduce(0) { $0 + $1.size * $1.itemSize } + } + } + + func reset() { + convolution.removeAll(); recurrent.removeAll() + keys.removeAll(); values.removeAll() + offset = 0 + pregate.reset() + recentRouting.removeAll(); prefetched.removeAll() + } + } + + let configuration: Configuration + private let weights: SafetensorsBundle + private let experts: ExpertLoader + private let lora: SafetensorsShard? + private let directory: URL + private let coreMLHead: CoreMLVocabularyHead? + private var coreMLHeadEnabled: Bool + private var coreMLHeadValidated = false + + init(directory: URL, configuration: Configuration = Configuration()) throws { + let index = directory.appendingPathComponent("model.safetensors.index.json") + guard FileManager.default.fileExists(atPath: index.path) else { + throw Failure.weightsMissing(directory) + } + self.directory = directory + self.configuration = configuration + self.weights = try SafetensorsBundle(directory: directory) + self.experts = try ExpertLoader( + directory: directory, layers: configuration.layers, + slots: Edge0Model35B.maximumBatch * configuration.activeExperts, + hidden: configuration.hidden, + moeIntermediate: configuration.moeIntermediate, + groupSize: configuration.groupSize) + self.experts.profile = self.profile + let round6LoRA = directory.appendingPathComponent("lora-round6.safetensors") + let officialLoRA = directory.appendingPathComponent("lora_edge0_35b.safetensors") + let legacyLoRA = directory.appendingPathComponent("lora.safetensors") + let loraURL: URL + if configuration.round6 { + loraURL = round6LoRA + } else if FileManager.default.fileExists(atPath: officialLoRA.path) { + loraURL = officialLoRA + } else { + loraURL = legacyLoRA + } + self.lora = try? SafetensorsShard(url: loraURL) + if configuration.round6 { + print("[round6] K=\(configuration.activeExperts) lora=\(self.lora == nil ? "missing" : "loaded")") + } + let compiledHead = directory.appendingPathComponent( + CoreMLVocabularyHead.directoryName, isDirectory: true) + if CommandLine.arguments.contains("--ane-35b"), + FileManager.default.fileExists(atPath: compiledHead.path) { + do { + self.coreMLHead = try CoreMLVocabularyHead(directory: directory) + self.coreMLHeadEnabled = true + print("[coreml-head] loaded with computeUnits=cpuAndNeuralEngine") + } catch { + self.coreMLHead = nil + self.coreMLHeadEnabled = false + print("[coreml-head] unavailable: \(error)") + } + } else { + self.coreMLHead = nil + self.coreMLHeadEnabled = false + print("[coreml-head] disabled or absent") + } + let free = (try? directory.resourceValues( + forKeys: [.volumeAvailableCapacityForImportantUsageKey]))? + .volumeAvailableCapacityForImportantUsage + print("[weights] free=\(free.map { String($0 / 1_048_576) } ?? "?")MiB") + + // Pregate weights are optional. Report an invalid present file so a fallback to + // the normal router is visible during diagnostics. + let pregateURL = PregateHeads.fileURL(in: directory, round6: configuration.round6) + if configuration.routingGuess == .trainedHeads, + FileManager.default.fileExists(atPath: pregateURL.path) { + do { + let heads = try PregateHeads( + url: pregateURL, hidden: configuration.hidden, + experts: RepackedExperts.expertsPerLayer) + self.pregate = heads + print("[pregate] heads loaded: owners \(heads.owners.first ?? -1)..\(heads.owners.last ?? -1) n=\(heads.owners.count)") + if let report = heads.selfCheck( + fixture: directory.appendingPathComponent("pregate-fixture.safetensors")) { + print(report) + } + } catch { + print("[pregate] present but unusable: \(error)") + self.pregate = nil + } + } else { + self.pregate = nil + } + } + + /// The trained routing heads, when the file is present. See `PregateHeads`. + private let pregate: PregateHeads? + + /// See `StepProfile` for why only the synchronous phases are timed. + let profile = StepProfile() + + var loraLoaded: Bool { lora != nil } + var availableLayers: Int { experts.availableLayers.count } + + // MARK: - Weights + + /// Selected tensors retained after first use to reduce repeated materialization. + private var retainedWeights: [String: MLXArray] = [:] + private var retainAttentionWeights = false + + /// First-turn prefill is latency-sensitive and sees every attention weight only + /// once, so retaining them there adds pressure without reuse. Decode immediately + /// reuses them; enable the cache at that boundary and keep it warm for later turns. + func beginDecodeCaching() { + retainAttentionWeights = true + } + + /// Prepare selected attention weights and issue a discarded token before generation + /// so initial Metal setup is outside the first reply's latency measurement. + func prepareForFirstTurn() throws { + beginDecodeCaching() + let started = Date() + var count = 0 + for name in weights.index.keys where + name.hasPrefix("language_model.lm_head.") + || name.contains(".linear_attn.") + || name.contains(".self_attn.") { + if weight(name) != nil { count += 1 } + } + print(String( + format: "[weights] first-turn cache tensors=%d %.2fs", + count, Date().timeIntervalSince(started))) + + let gpu = Date() + let scratch = State() + _ = try step(tokens: MLXArray([Int32(0)], [1, 1]), state: scratch) + profile.reset() + print(String(format: "[weights] gpu warm %.2fs", Date().timeIntervalSince(gpu))) + } + + /// Materialize a tensor from the mapped checkpoint. Retain selected tensors according + /// to the configured memory policy; other tensors are released after use. + private func weight(_ name: String) -> MLXArray? { + if let retained = retainedWeights[name] { return retained } + guard let entry = weights.entry(name), let pointer = weights.pointer(to: name) + else { return nil } + let raw = UnsafeRawBufferPointer(start: pointer, count: entry.byteCount) + profile.recordWeight(bytes: entry.byteCount) + let made = profile.measure(.weightCopy) { () -> MLXArray? in + switch entry.dtype { + case "U32", "I32": return MLXArray(raw, entry.shape, type: UInt32.self) + case "BF16": + return MLXArray(raw, entry.shape, type: UInt16.self).view(dtype: .bfloat16) + case "F32": return MLXArray(raw, entry.shape, type: Float.self) + default: return nil + } + } + let shouldRetain = configuration.retainAllResidentWeights + || (!coreMLHeadEnabled && name.hasPrefix("language_model.lm_head.")) + || (retainAttentionWeights + && (name.contains(".linear_attn.") || name.contains(".self_attn."))) + if shouldRetain, let made { + retainedWeights[name] = made + } + return made + } + + /// Apply the base quantized linear layer and its optional LoRA update. + private func linear(_ x: MLXArray, _ stem: String, bits: Int) -> MLXArray? { + guard let w = weight("\(stem).weight"), + let s = weight("\(stem).scales"), + let b = weight("\(stem).biases") + else { return nil } + let base = quantizedMatmul( + x.asType(.bfloat16), w, scales: s, biases: b, transpose: true, + groupSize: configuration.groupSize, bits: bits, mode: .affine) + + guard let prepared = preparedLoRA(stem) else { return base } + let delta = matmul(matmul(x.asType(.float16), prepared.a), prepared.b) + return base + delta.asType(base.dtype) + } + + /// LoRA matrices are converted, transposed, scaled, and cached on first use. + private struct PreparedLoRA { + let a: MLXArray + let b: MLXArray + } + private var loraCache: [String: PreparedLoRA] = [:] + + private func preparedLoRA(_ stem: String) -> PreparedLoRA? { + if let cached = loraCache[stem] { return cached } + guard let lora, + let aEntry = lora.entries["\(stem).lora_A"], + let bEntry = lora.entries["\(stem).lora_B"], + let aPointer = lora.pointer(to: "\(stem).lora_A"), + let bPointer = lora.pointer(to: "\(stem).lora_B") + else { return nil } + + func loadAdapter(_ entry: SafetensorsShard.Entry, + _ pointer: UnsafeRawPointer) -> MLXArray? { + let raw = UnsafeRawBufferPointer(start: pointer, count: entry.byteCount) + switch entry.dtype { + case "F16": return MLXArray(raw, entry.shape, type: Float16.self) + case "F32": return MLXArray(raw, entry.shape, type: Float.self) + default: return nil + } + } + guard let a = loadAdapter(aEntry, aPointer), + let bb = loadAdapter(bEntry, bPointer) + else { return nil } + // float16 throughout, as the Python implementation does. + let prepared = PreparedLoRA( + a: a.asType(.float16).transposed(), + b: (Float16(2.0) * bb.asType(.float16)).transposed()) + // Materialised here so the graph that produced them is not replayed per token. + eval(prepared.a, prepared.b) + loraCache[stem] = prepared + return prepared + } + + private func norm(_ x: MLXArray, _ name: String) -> MLXArray? { + guard let w = weight(name) else { return nil } + return MLXFast.rmsNorm(x.asType(.bfloat16), weight: w, eps: 1e-6) + } + + private func rmsNormalise(_ x: MLXArray) -> MLXArray { + let squares = mean(x.asType(.float32) * x.asType(.float32), axis: -1, keepDims: true) + return (x.asType(.float32) * rsqrt(squares + 1e-6)).asType(x.dtype) + } + + // MARK: - Layers + + /// Copy just the rows named, straight out of the mapping. + /// + /// `weight(_:)` builds an `MLXArray` over a whole tensor, and building one **copies** + /// — mlx-c routes it to `allocator::malloc` plus `std::copy`. For a table indexed one + /// row at a time that is the wrong shape by five orders of magnitude. + private func rows(_ name: String, _ indices: [Int32]) -> MLXArray? { + guard let entry = weights.entry(name), let base = weights.pointer(to: name), + entry.shape.count == 2, entry.shape[0] > 0 + else { return nil } + let columns = entry.shape[1] + let rowBytes = entry.byteCount / entry.shape[0] + var gathered = [UInt8](repeating: 0, count: rowBytes * indices.count) + gathered.withUnsafeMutableBytes { destination in + for (slot, row) in indices.enumerated() { + guard row >= 0, Int(row) < entry.shape[0] else { continue } + memcpy(destination.baseAddress!.advanced(by: slot * rowBytes), + base.advanced(by: Int(row) * rowBytes), rowBytes) + } + } + let shape = [indices.count, columns] + return gathered.withUnsafeBytes { raw -> MLXArray? in + switch entry.dtype { + case "U32", "I32": return MLXArray(raw, shape, type: UInt32.self) + case "BF16": return MLXArray(raw, shape, type: UInt16.self).view(dtype: .bfloat16) + case "F32": return MLXArray(raw, shape, type: Float.self) + default: return nil + } + } + } + + /// Gather only the embedding rows needed by the current token batch. + private func embed(_ tokens: MLXArray) -> MLXArray? { + let stem = "language_model.model.embed_tokens" + let ids = tokens.reshaped([-1]).asArray(Int32.self) + guard let w = rows("\(stem).weight", ids), + let s = rows("\(stem).scales", ids), + let b = rows("\(stem).biases", ids) + else { return nil } + // Dequantize only the gathered rows rather than the full embedding table. + let out = dequantized(w, scales: s, biases: b, + groupSize: configuration.groupSize, bits: configuration.bits, + mode: .affine) + return out.reshaped(tokens.shape + [out.dim(-1)]) + } + + /// Grouped-query attention with partial RoPE and an output gate. + /// + /// The gate is the second half of `q_proj`; RoPE covers 64 of each 256-dimensional + /// head and uses a base of 10,000,000. + private func fullAttention(_ x: MLXArray, layer: Int, state: State) -> MLXArray? { + let stem = "language_model.model.layers.\(layer).self_attn" + let heads = 16, kvHeads = 2, headDim = 256 + let (batch, length) = (x.dim(0), x.dim(1)) + + guard let qOut = linear(x, "\(stem).q_proj", bits: 4), + let kOut = linear(x, "\(stem).k_proj", bits: 4), + let vOut = linear(x, "\(stem).v_proj", bits: 4), + let qNorm = weight("\(stem).q_norm.weight"), + let kNorm = weight("\(stem).k_norm.weight") + else { return nil } + + let split = qOut.reshaped([batch, length, heads, 2 * headDim]) + var queries = split[.ellipsis, 0 ..< headDim] + let gate = split[.ellipsis, headDim ..< (2 * headDim)] + .reshaped([batch, length, heads * headDim]) + + var keys = kOut.reshaped([batch, length, kvHeads, headDim]) + var values = vOut.reshaped([batch, length, kvHeads, headDim]) + + queries = MLXFast.rmsNorm(queries, weight: qNorm, eps: 1e-6).transposed(0, 2, 1, 3) + keys = MLXFast.rmsNorm(keys, weight: kNorm, eps: 1e-6).transposed(0, 2, 1, 3) + values = values.transposed(0, 2, 1, 3) + + let offset = state.offset + func rope(_ a: MLXArray) -> MLXArray { + MLXFast.RoPE(a, dimensions: 64, traditional: false, + base: 10_000_000, scale: 1, offset: offset) + } + queries = rope(queries) + keys = rope(keys) + + if let pastKeys = state.keys[layer], let pastValues = state.values[layer] { + keys = concatenated([pastKeys, keys], axis: 2) + values = concatenated([pastValues, values], axis: 2) + } + state.keys[layer] = keys + state.values[layer] = values + + // A single query attending to everything before it needs no mask; passing a + // square one sized for the prompt is how a working prefill breaks decode. + let total = keys.dim(2) + var mask: MLXArray? = nil + if length > 1 { + var entries = [Float](repeating: 0, count: length * total) + for row in 0 ..< length { + for column in (offset + row + 1) ..< total { + entries[row * total + column] = -Float.greatestFiniteMagnitude + } + } + mask = MLXArray(entries, [length, total]).asType(queries.dtype) + } + + let attended = MLXFast.scaledDotProductAttention( + queries: queries, keys: keys, values: values, scale: 1.0 / 16.0, mask: mask) + let merged = attended.transposed(0, 2, 1, 3) + .reshaped([batch, length, heads * headDim]) + return linear(merged * sigmoid(gate), "\(stem).o_proj", bits: 4) + } + + /// Gated DeltaNet: depthwise causal convolution, then a gated delta-rule recurrence + /// over a float32 state. + /// + /// `q` and `k` carry 16 heads against `v`'s 32 and are matched by **consecutive** + /// duplication — h0, h0, h1, h1 — not by tiling. Both run; one is correct. + private func linearAttention(_ x: MLXArray, layer: Int, state: State) -> MLXArray? { + let stem = "language_model.model.layers.\(layer).linear_attn" + let kHeads = 16, vHeads = 32, headDim = 128 + let keyDim = kHeads * headDim, valueDim = vHeads * headDim + let convDim = keyDim * 2 + valueDim, kernel = 4 + let (batch, length) = (x.dim(0), x.dim(1)) + + guard let qkv = linear(x, "\(stem).in_proj_qkv", bits: 4), + let zRaw = linear(x, "\(stem).in_proj_z", bits: 4), + let aRaw = linear(x, "\(stem).in_proj_a", bits: 4), + let bRaw = linear(x, "\(stem).in_proj_b", bits: 4), + let convWeight = weight("\(stem).conv1d.weight"), + let aLog = weight("\(stem).A_log"), + let dtBias = weight("\(stem).dt_bias"), + let normWeight = weight("\(stem).norm.weight") + else { return nil } + + let history = state.convolution[layer]?.asType(qkv.dtype) + ?? MLXArray.zeros([batch, kernel - 1, convDim]).asType(qkv.dtype) + let tail = concatenated([history, qkv], axis: 1) + + // `conv1d` accumulates internally in float32. Preserve that precision for + // numerically consistent recurrent state updates. + let convolved = conv1d(tail, convWeight, groups: convDim) + let activated = convolved * sigmoid(convolved) + state.convolution[layer] = tail[0..., (tail.dim(1) - (kernel - 1))...] + + var q = activated[.ellipsis, 0 ..< keyDim].reshaped([batch, length, kHeads, headDim]) + var k = activated[.ellipsis, keyDim ..< (2 * keyDim)] + .reshaped([batch, length, kHeads, headDim]) + let v = activated[.ellipsis, (2 * keyDim) ..< convDim] + .reshaped([batch, length, vHeads, headDim]) + + // q is scaled by Dk⁻¹ and k by Dk^-0.5 — not the same factor. + let inverseScale = Float(pow(Double(headDim), -0.5)) + q = rmsNormalise(q) * (inverseScale * inverseScale) + k = rmsNormalise(k) * inverseScale + q = repeated(q, count: vHeads / kHeads, axis: -2) + k = repeated(k, count: vHeads / kHeads, axis: -2) + + let beta = sigmoid(bRaw.asType(.float32)) + let shifted = aRaw.asType(.float32) + dtBias.asType(.float32) + let g = MLX.exp(-MLX.exp(aLog.asType(.float32)) * MLX.log(MLX.exp(shifted) + 1)) + + var recurrent = state.recurrent[layer]?.asType(.float32) + ?? MLXArray.zeros([batch, vHeads, headDim, headDim], dtype: .float32) + var outputs: [MLXArray] = [] + // Use matrix contractions for the recurrent updates to avoid materializing large + // broadcast-multiply intermediates: + // + // (recurrent * kt[..., None, :]).sum(-1) 2 MiB temporary, reduced to 16 KB + // matmul(recurrent, kt[..., None]) 16 KB, straight out of a GEMV + // + // GEMV accumulation order can differ from elementwise reduction in the last bits. + for t in 0 ..< length { + let qt = q[0..., t].asType(.float32).expandedDimensions(axis: -1) + let kt = k[0..., t].asType(.float32).expandedDimensions(axis: -1) + let vt = v[0..., t].asType(.float32) + let decay = g[0..., t].expandedDimensions(axis: -1).expandedDimensions(axis: -1) + + recurrent = recurrent * decay + let memory = matmul(recurrent, kt).squeezed(axis: -1) + let delta = (vt - memory) * beta[0..., t].expandedDimensions(axis: -1) + // The outer product has to materialise 2 MiB whichever way it is written, so + // this one stays as it was. + recurrent = recurrent + + kt.swappedAxes(-1, -2) * delta.expandedDimensions(axis: -1) + outputs.append(matmul(recurrent, qt).squeezed(axis: -1)) + } + state.recurrent[layer] = recurrent + + let stacked = MLX.stacked(outputs, axis: 1).asType(x.dtype) + let z = zRaw.reshaped([batch, length, vHeads, headDim]) + // RMSNormGated, in float32 as trained. + let normalised = MLXFast.rmsNorm(stacked, weight: normWeight, eps: 1e-6) + .asType(.float32) + let gated = (sigmoid(z.asType(.float32)) * z.asType(.float32) * normalised) + .asType(x.dtype) + return linear(gated.reshaped([batch, length, valueDim]), "\(stem).out_proj", bits: 4) + } + + /// Router, K streamed experts, and the shared expert. + /// + /// Apply softmax over all expert logits before top-k selection, then renormalize the + /// selected scores. + /// Routed experts for **one or more** tokens. + /// + /// The single-token form this replaces treated the K experts as the batch axis, every + /// one of them seeing the same input. With T tokens there are T*K (token, expert) + /// pairs and each needs its own input row, so the batch axis becomes T*K and the + /// expert for each row is named by `rhsIndices`. + /// + /// For a batch, read the union of experts selected by its tokens. The union may be + /// smaller than T*K when tokens share expert selections. + private func moe(_ x: MLXArray, layer: Int, state: State) -> MLXArray? { + let prefix = "language_model.model.layers.\(layer).mlp" + let k = configuration.activeExperts + let tokens = x.dim(1) + + if tokens >= Self.fullLayerPrefillMinimum { + return moeFullLayer(x, layer: layer, state: state, prefix: prefix, k: k, tokens: tokens) + } + + let ids: [Int] + var scores: MLXArray + + // Predictions replace the normal gate during single-token decode. Prefill uses + // the normal gate because its multi-token routing shape differs. + if tokens == 1, let predicted = state.pregate.predicted[layer], + let predictedScores = state.pregate.scores[layer], predicted.count == k { + ids = predicted + // Rebuilt here, from k floats. The alternative — keeping the slice of the + // staging array — would hold that array and its graph across the step. + scores = MLXArray(predictedScores, [1, 1, k]) + } else { + guard let logits = linear(x, "\(prefix).gate", bits: 8) else { return nil } + let gates = softmax(logits.asType(.float32), axis: -1) + let chosen = argPartition(-gates, kth: k - 1, axis: -1)[.ellipsis, 0 ..< k] + scores = takeAlong(gates, chosen, axis: -1) + scores = scores / scores.sum(axis: -1, keepDims: true) + + // Expert IDs must reach the CPU before their file reads can be issued. + let kind: StepProfile.Phase = + configuration.isLinear(layer) ? .syncLinear : .syncFull + ids = profile.measure(kind) { () -> [Int] in + eval(chosen, scores) + return chosen.reshaped([-1]).asArray(Int32.self).map { Int($0) } + } + } + guard ids.count == tokens * k else { return nil } + + // Captured for this layer's own head, which runs at the step boundary and + // predicts for `layer + 1`. The last position is the right one for both: a + // prefill's final position is what the first decode step continues from. + // + // Store the routes actually used, including predictions, as inputs for the next + // pregate evaluation. Tensor reshaping is deferred until after graph evaluation. + if let pregate, pregate.owners.contains(layer) { + state.pregate.moeInput[layer] = x + state.pregate.executed[layer] = ids.suffix(k).map { Int32($0) } + } + + // Distinct experts in first-appearance order, and where each pair's expert landed. + var rowOf: [Int: Int] = [:] + var union: [Int] = [] + for id in ids where rowOf[id] == nil { + rowOf[id] = union.count + union.append(id) + } + profile.recordRouting(union: union.count, pairs: ids.count) + + // How much of what this layer needs was already asked for. **This is the number + // that bounds the whole scheme** — the read can only be hidden for blocks the + // guess anticipated — and it is free to collect, because both sets are already + // on the CPU. Recorded before the routing is overwritten below. + if configuration.routingGuess != .none { + let guessed = state.prefetched[layer] ?? [] + profile.recordPrefetch( + hits: union.count { guessed.contains($0) }, total: union.count) + } + // What executed, for the next step to guess from. The union rather than the + // last position's pair: during prefill that is strictly more informative, and + // during decode they are the same thing. + var history = state.recentRouting[layer] ?? [] + history.insert(union, at: 0) + state.recentRouting[layer] = Array(history.prefix(configuration.routingLookback)) + + guard let pool = try? experts.load(layer: layer, experts: union), + let rows = try? ids.map({ id -> Int32 in + guard let row = rowOf[id] else { throw Failure.layerFailed(layer) } + return Int32(row) + }) + else { return nil } + + // **The single-token form is kept separate, and it is not premature.** Generalising + // this for batching added a `reshaped` on top of a broadcast — a stride-0 view that + // a reshape has to materialise — and decode runs this on every token of every + // layer. Batching is used only by prefill; making decode pay for it was a + // regression, not a simplification. + let expanded = x.reshaped([tokens, 1, configuration.hidden]).asType(.bfloat16) + let inputs = tokens == 1 + ? broadcast(expanded, to: [k, 1, configuration.hidden]) + : broadcast(expanded, to: [tokens, k, configuration.hidden]) + .reshaped([tokens * k, 1, configuration.hidden]) + let indices = MLXArray(rows) + + func project(_ input: MLXArray, _ offset: Int) -> MLXArray { + gatherQuantizedMM( + input, pool[offset], scales: pool[offset + 1], biases: pool[offset + 2], + rhsIndices: indices, transpose: true, + groupSize: configuration.groupSize, bits: configuration.bits, + // **An assertion about the input, not a request**, so it is true only + // when it is true. With one token the union is built in first-appearance + // order from k distinct experts, making the rows exactly `0 ..< k` and + // ascending; with more, rows repeat and are unordered, and claiming + // otherwise would silently gather the wrong experts. + mode: .affine, sortedIndices: tokens == 1) + } + let gate = project(inputs, ExpertLoader.gate) + let up = project(inputs, ExpertLoader.up) + let routed = project(gate * sigmoid(gate) * up, ExpertLoader.down) + let weighted = routed + * scores.reshaped([tokens * k, 1, 1]).asType(routed.dtype) + let combined = weighted + .reshaped([tokens, k, configuration.hidden]) + .sum(axis: 1) + + // Compute the shared expert after routed expert loading to avoid competing for + // memory bandwidth during the reads. + guard let sharedGate = linear(x, "\(prefix).shared_expert_gate", bits: 8), + let sg = linear(x, "\(prefix).shared_expert.gate_proj", bits: 4), + let su = linear(x, "\(prefix).shared_expert.up_proj", bits: 4), + let shared = linear(sg * sigmoid(sg) * su, + "\(prefix).shared_expert.down_proj", bits: 4) + else { return nil } + return combined.reshaped(shared.shape) + sigmoid(sharedGate) * shared + } + + /// Prefill for a long chunk: the layer file is read once, and expert ids + /// index that stack directly. The ids still come back to the CPU afterwards + /// so the pregate feature sees what actually ran; the read no longer waits + /// on that sync. + private func moeFullLayer(_ x: MLXArray, layer: Int, state: State, + prefix: String, k: Int, tokens: Int) -> MLXArray? { + guard let pool = try? experts.loadFullLayer(layer: layer), + let logits = linear(x, "\(prefix).gate", bits: 8) + else { return nil } + let gates = softmax(logits.asType(.float32), axis: -1) + let chosen = argPartition(-gates, kth: k - 1, axis: -1)[.ellipsis, 0 ..< k] + var scores = takeAlong(gates, chosen, axis: -1) + scores = scores / scores.sum(axis: -1, keepDims: true) + + let kind: StepProfile.Phase = + configuration.isLinear(layer) ? .syncLinear : .syncFull + let ids: [Int] = profile.measure(kind) { + eval(chosen, scores) + return chosen.reshaped([-1]).asArray(Int32.self).map { Int($0) } + } + guard ids.count == tokens * k else { return nil } + if let pregate, pregate.owners.contains(layer) { + state.pregate.moeInput[layer] = x + state.pregate.executed[layer] = ids.suffix(k).map { Int32($0) } + } + var rowOf: [Int: Int] = [:] + var union: [Int] = [] + for id in ids where rowOf[id] == nil { + rowOf[id] = union.count + union.append(id) + } + profile.recordRouting(union: union.count, pairs: ids.count) + var history = state.recentRouting[layer] ?? [] + history.insert(union, at: 0) + state.recentRouting[layer] = Array(history.prefix(configuration.routingLookback)) + + let expanded = x.reshaped([tokens, 1, configuration.hidden]).asType(.bfloat16) + let inputs = broadcast(expanded, to: [tokens, k, configuration.hidden]) + .reshaped([tokens * k, 1, configuration.hidden]) + // Row `e` of the full stack is expert `e`, so the router index is the gather index. + let indices = chosen.reshaped([tokens * k]) + func project(_ input: MLXArray, _ offset: Int) -> MLXArray { + gatherQuantizedMM( + input, pool[offset], scales: pool[offset + 1], biases: pool[offset + 2], + rhsIndices: indices, transpose: true, + groupSize: configuration.groupSize, bits: configuration.bits, + mode: .affine, sortedIndices: false) + } + let gate = project(inputs, ExpertLoader.gate) + let up = project(inputs, ExpertLoader.up) + let routed = project(gate * sigmoid(gate) * up, ExpertLoader.down) + let weighted = routed * scores.reshaped([tokens * k, 1, 1]).asType(routed.dtype) + let combined = weighted.reshaped([tokens, k, configuration.hidden]).sum(axis: 1) + guard let sharedGate = linear(x, "\(prefix).shared_expert_gate", bits: 8), + let sg = linear(x, "\(prefix).shared_expert.gate_proj", bits: 4), + let su = linear(x, "\(prefix).shared_expert.up_proj", bits: 4), + let shared = linear(sg * sigmoid(sg) * su, + "\(prefix).shared_expert.down_proj", bits: 4) + else { return nil } + return combined.reshaped(shared.shape) + sigmoid(sharedGate) * shared + } + + /// Run all routing heads as one batch and return the expert IDs for the next step. + /// + /// Returns the expert set per consuming layer, which is what a prefetch needs. + @discardableResult + private func stagePregate(state: State) -> [Int: [Int]] { + guard let pregate else { return [:] } + let k = configuration.activeExperts + + var inputs: [MLXArray] = [] + var executed: [[Int32]] = [], previous: [[Int32]] = [] + inputs.reserveCapacity(pregate.owners.count) + for owner in pregate.owners { + // A pruned or otherwise absent owner would desynchronise the stack from + // `owners`, so the whole staging is skipped rather than silently shifted. + // + // Skipping has to *clear* the predictions, not just decline to write new + // ones. `advance()` deliberately keeps `predicted` alive across the step + // boundary — that is how a prediction made at step N reaches step N+1 — so + // leaving it untouched here would hand step N+1 the routing computed for + // step N, one token stale, with nothing in any log to say so. + guard let input = state.pregate.moeInput[owner], + let ran = state.pregate.executed[owner] + else { + state.pregate.predicted = [:] + state.pregate.scores = [:] + state.pregate.advance() + return [:] + } + // Sliced here rather than at capture time: by now `eval(h)` has run, so this + // reshapes a materialised array instead of adding a node to a live graph. + // The last position is the right one for a prefill chunk too — it is what + // the first decode step continues from. + let vector = input.ndim == 1 + ? input + : (input.dim(1) == 1 ? input.reshaped([-1]) : input[0, -1]) + inputs.append(vector) + executed.append(ran) + // Use a zero row when no previous-token routing is available. + previous.append(state.pregate.previous[owner] ?? []) + } + + let logits = pregate.logits( + inputs: stacked(inputs, axis: 0), + executed: pregate.oneHotRows(executed), + previous: pregate.oneHotRows(previous)) + + // Apply the same softmax, top-k selection, and score normalization as the router. + let gates = softmax(logits, axis: -1) + let chosen = argPartition(-gates, kth: k - 1, axis: -1)[.ellipsis, 0 ..< k] + var weights = takeAlong(gates, chosen, axis: -1) + weights = weights / weights.sum(axis: -1, keepDims: true) + + // The one barrier, and it brings back both halves. Pulling the weights across + // too costs 66 floats and means nothing MLX-shaped is carried into the next + // step — a view into `weights` would have kept the array, and its graph, alive + // for every one of the thirty-three layers that held one. + let (flat, weighting) = profile.measure(.pregateStage) { () -> ([Int32], [Float]) in + eval(chosen, weights) + return (chosen.reshaped([-1]).asArray(Int32.self), + weights.reshaped([-1]).asArray(Float.self)) + } + + var sets: [Int: [Int]] = [:] + for (index, consumer) in pregate.consumers.enumerated() { + let range = (index * k) ..< ((index + 1) * k) + let slice = flat[range].map { Int($0) } + state.pregate.predicted[consumer] = slice + state.pregate.scores[consumer] = Array(weighting[range]) + sets[consumer] = Array(Set(slice)).sorted() + } + + state.pregate.advance() + return sets + } + + // MARK: - Forward + + /// Run one model step and evaluate the resulting graph once at the end. + func step(tokens: MLXArray, state: State) throws -> MLXArray { + try step(tokens: tokens, state: state, needsLogits: true)! + } + + /// One step, with the option of **not** projecting to the vocabulary. + /// + /// During prefill, intermediate positions do not need vocabulary logits. This option + /// skips that projection when the result will not be consumed. + func step(tokens: MLXArray, state: State, needsLogits: Bool) throws -> MLXArray? { + guard var h = embed(tokens) else { throw Failure.tensorMissing("embed_tokens") } + // Checkpoint weights use bfloat16; keep activations in the same format. + h = h.asType(.bfloat16) + + let batchedPrefill = tokens.dim(1) >= Self.fullLayerPrefillMinimum + for layer in 0 ..< configuration.layers { + if batchedPrefill && layer % 10 == 0 { + print("[35b layer] \(layer) T=\(tokens.dim(1))") + } + let stem = "language_model.model.layers.\(layer)" + guard let normed = norm(h, "\(stem).input_layernorm.weight") else { + throw Failure.layerFailed(layer) + } + let attended = configuration.isLinear(layer) + ? linearAttention(normed, layer: layer, state: state) + : fullAttention(normed, layer: layer, state: state) + guard let attended else { throw Failure.layerFailed(layer) } + h = h + attended + + guard let post = norm(h, "\(stem).post_attention_layernorm.weight"), + let mixed = moe(post, layer: layer, state: state) + else { throw Failure.layerFailed(layer) } + h = h + mixed + // A full-layer expert stack is 432 MiB. Evaluating here drops it + // before the next layer's read; holding all forty would be 17 GB. + if batchedPrefill { + eval(h) + if let held = state.pregate.moeInput[layer] { + let last = held[0, -1] + eval(last) + state.pregate.moeInput[layer] = last + } + } + } + eval(h) + + // Run routing predictions after graph evaluation; they consume hidden states + // produced by this step and provide expert IDs for the next step. + switch configuration.routingGuess { + case .none: + break + case .trainedHeads: + if tokens.dim(1) > 1 && !needsLogits { + // Batched prefill never consumes the one-token-ahead head predictions. + // Reuse the real routing as a read hint and avoid loading/running 132 MiB + // of heads until the final prefill batch, whose prediction is consumed + // by the first decode step. + state.prefetched = state.recentRouting.mapValues { Array(Set($0.joined())) } + for (layer, chosen) in state.prefetched { + experts.prefetch(layer: layer, experts: chosen) + } + break + } + state.prefetched = [:] + for (layer, chosen) in stagePregate(state: state) { + state.prefetched[layer] = chosen + experts.prefetch(layer: layer, experts: chosen) + } + // The official heads cover consumers 7...39. Preserve their exact routing, + // but prefetch the uncovered first layers from their last real routing. + // This is advice only: misses still run the router and read the right files. + for layer in 0 ..< configuration.layers where state.prefetched[layer] == nil { + guard let chosen = state.recentRouting[layer]?.first else { continue } + state.prefetched[layer] = chosen + experts.prefetch(layer: layer, experts: chosen) + } + case .recentRouting: + // Snapshotted, not read live: `recentRouting` is rewritten layer by layer as + // the next step runs, so scoring the hit rate against it would compare a + // layer's routing with its own and report a meaningless 100%. + state.prefetched = state.recentRouting.mapValues { Array(Set($0.joined())) } + for (layer, chosen) in state.prefetched { + experts.prefetch(layer: layer, experts: chosen) + } + } + + state.offset += tokens.dim(1) + profile.countStep() + guard needsLogits else { return nil } + guard let final = norm(h, "language_model.model.norm.weight") else { + throw Failure.tensorMissing("model.norm") + } + if coreMLHeadEnabled { + // Core ML consumes only the final position. This is 4 KiB instead of the + // 248,320-element logits vector, and avoids retaining MLX's 286 MiB head. + let last = final[0, -1].reshaped([1, configuration.hidden]) + eval(last) + return last + } + guard let logits = linear(final, "language_model.lm_head", bits: 4) else { + throw Failure.tensorMissing("lm_head") + } + return logits + } + + /// Argmax over the last position's logits. + func greedy(_ output: MLXArray) async throws -> Int32 { + if coreMLHeadEnabled, let coreMLHead, + output.dim(output.ndim - 1) == configuration.hidden { + let coreMLToken = try await profile.measure(.vocabularyHead) { + try await coreMLHead.predict(output) + } + if !coreMLHeadValidated { + guard let logits = linear(output, "language_model.lm_head", bits: 4) else { + throw Failure.tensorMissing("lm_head") + } + let mlxToken = Self.mlxGreedy(logits) + if mlxToken != coreMLToken { + coreMLHeadEnabled = false + print("[coreml-head] parity FAIL coreml=\(coreMLToken) mlx=\(mlxToken); falling back") + return mlxToken + } + coreMLHeadValidated = true + retainedWeights = retainedWeights.filter { + !$0.key.hasPrefix("language_model.lm_head.") + } + print("[coreml-head] parity PASS token=\(coreMLToken)") + } + return coreMLToken + } + return Self.mlxGreedy(output) + } + + private static func mlxGreedy(_ logits: MLXArray) -> Int32 { + // `-1`, not `0`: with a batched forward the interesting logits are the last + // position's. Identical for a single token, which is why this was never wrong. + let final = logits.ndim == 2 ? logits[0] : logits[0, -1] + let choice = argMax(final.asType(.float32), axis: -1) + eval(choice) + return choice.item(Int32.self) + } + + /// Greedy decode. Prompt and generation take the same path — one token at a time — + /// which costs a little on prefill and buys a single code path. + func generate(promptIds: [Int32], stop: StopCondition, + state: State? = nil, + onToken: @escaping (Int32) -> Bool) throws -> [Int32] + { + let state = state ?? State() + var produced: [Int32] = [] + var next: Int32 = 0 + + for (index, token) in promptIds.enumerated() { + let logits = try step(tokens: MLXArray([token], [1, 1]), state: state) + if index == promptIds.count - 1 { + let choice = argMax(logits[0, 0].asType(.float32), axis: -1) + eval(choice) + next = choice.item(Int32.self) + } + } + + while true { + switch stop.evaluate(token: next, produced: produced.count) { + case .hitEndToken, .hitLimit: + return produced + case .keepGoing: + break + } + produced.append(next) + // The callback decides whether to continue, so cancellation is the caller's + // to express rather than something this loop has to be told about. + guard onToken(next) else { return produced } + + let logits = try step(tokens: MLXArray([next], [1, 1]), state: state) + let choice = argMax(logits[0, 0].asType(.float32), axis: -1) + eval(choice) + next = choice.item(Int32.self) + } + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift b/ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift new file mode 100644 index 0000000..1687f20 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift @@ -0,0 +1,241 @@ +import Foundation +import MLX + +/// Loads routed experts from the per-layer files and builds the stacked tensors used by +/// `gatherQuantizedMM`. Reads for the selected experts are issued concurrently. +final class ExpertLoader: @unchecked Sendable { + + enum Failure: LocalizedError { + case noLayers(URL) + case unknownLayer(Int, available: Int) + case expertOutOfRange(Int) + case wrongSlotCount(asked: Int, configured: Int) + case readFailed(layer: Int, expert: Int) + + var errorDescription: String? { + switch self { + case .noLayers(let url): + return "no intact layer files in \(url.lastPathComponent)" + case .unknownLayer(let layer, let available): + return "layer \(layer) is not loaded (\(available) layers available)" + case .expertOutOfRange(let expert): + return "expert \(expert) outside 0..<\(RepackedExperts.expertsPerLayer)" + case .wrongSlotCount(let asked, let configured): + return "asked for \(asked) experts, loader is built for \(configured)" + case .readFailed(let layer, let expert): + return "read failed for layer \(layer) expert \(expert)" + } + } + } + + /// Where each tensor sits in the stacked pool handed to `gatherQuantizedMM`: + /// weight, scales, biases for gate, then up, then down. + static let gate = 0, up = 3, down = 6 + + private let experts: RepackedExperts + private let slots: Int + private var buffers: [UnsafeMutableRawPointer] = [] + + + private let shapes: [(weight: [Int], scale: [Int])] + + /// Set by the model so reads and array construction land in the same accounting as + /// the router barrier. Optional so the loader stays usable without it. + var profile: StepProfile? + + var availableLayers: [Int] { experts.availableLayers } + var rejectedFiles: [(layer: Int, bytes: Int)] { experts.rejected } + + init(directory: URL, layers: Int, slots: Int, hidden: Int, + moeIntermediate: Int, groupSize: Int) throws + { + self.experts = RepackedExperts(directory: directory, expectedLayers: layers) + guard !experts.availableLayers.isEmpty else { throw Failure.noLayers(directory) } + + self.slots = slots + + self.shapes = [ + (weight: [moeIntermediate, hidden / 8], scale: [moeIntermediate, hidden / groupSize]), + (weight: [moeIntermediate, hidden / 8], scale: [moeIntermediate, hidden / groupSize]), + (weight: [hidden, moeIntermediate / 8], scale: [hidden, moeIntermediate / groupSize]), + ] + } + + /// Nine tensors: gate weight/scales/biases, then up, then down. + /// + /// Throws rather than returning something plausible. A silently wrong expert + /// produces fluent nonsense several layers later, which is far more expensive to + /// find than a thrown error here. + deinit { + buffers.forEach { $0.deallocate() } + } + + /// Ask the kernel to prefetch the selected expert blocks into the file cache. + func prefetch(layer: Int, experts chosen: [Int]) { + for expert in chosen where expert >= 0 && expert < RepackedExperts.expertsPerLayer { + experts.prefetch(layer: layer, expert: expert) + } + } + + /// Load **1 to `slots`** distinct experts. + /// + /// Was fixed at exactly `slots`, which is what a single token needs at K=2. A batch + /// of T tokens routes to the *union* of their choices — anywhere from K (all tokens + /// agree) to T*K (none do) — so `slots` became the ceiling rather than the count. + func load(layer: Int, experts chosen: [Int]) throws -> [MLXArray] { + guard !chosen.isEmpty, chosen.count <= slots else { + throw Failure.wrongSlotCount(asked: chosen.count, configured: slots) + } + let count = chosen.count + guard experts.availableLayers.contains(layer) else { + throw Failure.unknownLayer(layer, available: experts.availableLayers.count) + } + for expert in chosen where expert < 0 || expert >= RepackedExperts.expertsPerLayer { + throw Failure.expertOutOfRange(expert) + } + + // A larger prefill batch needs a larger *possible* expert union, but reserving + // T*K blocks eagerly makes the worst case permanent. Grow to the union we have + // actually observed and reuse that high-water mark on later layers/turns. + while buffers.count < count { + buffers.append(UnsafeMutableRawPointer.allocate( + byteCount: RepackedExperts.blockBytes, + alignment: RepackedExperts.pageSize)) + } + + // Read each expert block into reusable aligned buffers, then stack its tensors. + let sources = Array(buffers.prefix(count)) + let missing = Array(0 ..< count) + + var failures = [Bool](repeating: false, count: missing.count) + let buffers = self.buffers + let experts = self.experts + let read = { + failures.withUnsafeMutableBufferPointer { flags in + let flagged = flags + let slots = missing + DispatchQueue.concurrentPerform(iterations: slots.count) { index in + let slot = slots[index] + flagged[index] = !experts.read( + layer: layer, expert: chosen[slot], into: buffers[slot]) + } + } + } + if let profile { profile.measure(.expertRead, read) } else { read() } + profile?.recordRead(bytes: missing.count * RepackedExperts.blockBytes, + issued: missing.count) + if let index = failures.firstIndex(of: true) { + throw Failure.readFailed(layer: layer, expert: chosen[missing[index]]) + } + + // `stack` is where `MLXArray(raw:)` runs, and that is `malloc` + `std::copy` + // in mlx-c — eager, so timing it needs no forced evaluation. + var pool: [MLXArray] = [] + pool.reserveCapacity(9) + let shapes = self.shapes + let build: () -> Void = { + for (projection, parts) in [ + (0, (RepackedExperts.Part.gateWeight, RepackedExperts.Part.gateScales, + RepackedExperts.Part.gateBiases)), + (1, (.upWeight, .upScales, .upBiases)), + (2, (.downWeight, .downScales, .downBiases)), + ] { + let shape = shapes[projection] + pool.append(self.stack(parts.0, shape.weight, bf16: false, from: sources)) + pool.append(self.stack(parts.1, shape.scale, bf16: true, from: sources)) + pool.append(self.stack(parts.2, shape.scale, bf16: true, from: sources)) + } + } + if let profile { profile.measure(.expertStack, build) } else { build() } + return pool + } + + /// Every expert in the layer, stacked as `[256, …]` in expert-id order. + /// + /// The file is one sequential 432 MiB read. Parts are then packed into the + /// nine tensors `gatherQuantizedMM` already understands, so a prefill chunk + /// can gather by expert id without a per-expert `pread` or a CPU copy of + /// the index list. The file buffer is released before return; MLX has + /// copied the nine tensors. Callers must `eval` and drop those tensors + /// before the next layer, or forty of them become 17 GB. + func loadFullLayer(layer: Int) throws -> [MLXArray] { + guard experts.availableLayers.contains(layer) else { + throw Failure.unknownLayer(layer, available: experts.availableLayers.count) + } + let file = UnsafeMutableRawPointer.allocate( + byteCount: RepackedExperts.layerBytes, alignment: RepackedExperts.pageSize) + defer { file.deallocate() } + + let read: () -> Bool = { self.experts.readLayer(layer, into: file) } + let ok = profile?.measure(.expertRead, read) ?? read() + guard ok else { throw Failure.readFailed(layer: layer, expert: -1) } + profile?.recordRead(bytes: RepackedExperts.layerBytes, issued: 1) + + let shapes = self.shapes + let parts: [(RepackedExperts.Part, RepackedExperts.Part, RepackedExperts.Part)] = [ + (.gateWeight, .gateScales, .gateBiases), + (.upWeight, .upScales, .upBiases), + (.downWeight, .downScales, .downBiases), + ] + var pool: [MLXArray] = [] + pool.reserveCapacity(9) + let build: () -> Void = { + for (projection, triple) in parts.enumerated() { + let shape = shapes[projection] + for (part, bf16, partShape) in [ + (triple.0, false, shape.weight), + (triple.1, true, shape.scale), + (triple.2, true, shape.scale), + ] { + pool.append(self.stackedPart( + part, partShape, bf16: bf16, file: file)) + } + } + } + if let profile { profile.measure(.expertStack, build) } else { build() } + return pool + } + + /// Expert id `e` lives at byte `e * block` in the layer file. Gather one + /// part from all 256 blocks into a contiguous `[256, …]` tensor. + private func stackedPart(_ part: RepackedExperts.Part, _ shape: [Int], + bf16: Bool, file: UnsafeMutableRawPointer) -> MLXArray { + let count = RepackedExperts.expertsPerLayer + let bytes = part.byteCount + let packed = UnsafeMutableRawPointer.allocate( + byteCount: count * bytes, alignment: 64) + defer { packed.deallocate() } + for expert in 0 ..< count { + packed.advanced(by: expert * bytes).copyMemory( + from: file.advanced(by: expert * RepackedExperts.blockBytes + part.offset), + byteCount: bytes) + } + let raw = UnsafeRawBufferPointer(start: packed, count: count * bytes) + let stackedShape = [count] + shape + return bf16 + ? MLXArray(raw, stackedShape, type: UInt16.self).view(dtype: .bfloat16) + : MLXArray(raw, stackedShape, type: UInt32.self) + } + + /// One part across all K slots, stacked along a new leading axis. + /// + /// bf16 is reached by reinterpreting the bits, never by converting values. An early + /// version widened each half in a Swift loop — roughly 200,000 iterations per layer + /// — and that single artifact produced the spike's original 0.5 tok/s figure. + private func stack(_ part: RepackedExperts.Part, _ shape: [Int], bf16: Bool, + from sources: [UnsafeMutableRawPointer]) -> MLXArray + { + // Sources are the staging buffers `load` has just filled. `MLXArray(raw:)` + // copies eagerly, so they are free to be overwritten the moment this returns. + let slices = sources.map { buffer -> MLXArray in + let raw = UnsafeRawBufferPointer( + start: buffer.advanced(by: part.offset), count: part.byteCount) + return bf16 + ? MLXArray(raw, shape, type: UInt16.self).view(dtype: .bfloat16) + : MLXArray(raw, shape, type: UInt32.self) + } + return slices.count == 1 + ? slices[0].expandedDimensions(axis: 0) + : MLX.stacked(slices, axis: 0) + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift b/ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift new file mode 100644 index 0000000..25b980e --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift @@ -0,0 +1,192 @@ +import Foundation +import MLX + +/// Decode state captured after a fixed run of prompt tokens, so those tokens are never +/// processed twice. +/// +/// **This is the single largest win available on time-to-first-character.** A measured +/// prefill runs at 258 ms/token, and of a 96-token prompt the great majority is the +/// system instructions — byte-identical on every request, in every conversation, for the +/// life of the install. Recomputing it is the bulk of a ~25 second wait. +/// +/// ## Why it is unusually cheap here +/// +/// Measured on device: **63 MiB**, against a 2.31 GB peak and a ~3.5 GB jetsam line. +/// +/// Sixty of those MiB are the thirty DeltaNet layers' recurrent state, whose size **does +/// not depend on how many tokens produced it** — caching a 96-token prefix and a +/// 10,000-token prefix cost the same. Only the ten full-attention layers contribute a +/// per-token term, at 2 KiB each. +/// +/// The property that makes those thirty layers awkward elsewhere — the state is a fixed +/// matrix updated in place, so dropping history does not shrink it and `replaceHistory` +/// cannot prune cheaply — is exactly the property that makes them nearly free to cache. +/// A plain forty-layer transformer of these dimensions would need 80 KiB per token. +/// +/// ## Why the prefix is verified rather than assumed +/// +/// The obvious implementation trusts that encoding the system message alone yields the +/// same tokens as encoding it as the head of the full prompt. Byte-level BPE gives no +/// such guarantee in general: merges are free to cross a boundary the renderer considers +/// structural. So the cached token run is stored alongside the state and checked against +/// the real prompt on every use. A mismatch costs an array comparison and falls back to +/// prefilling normally — it can never produce a wrong answer. +/// +/// This also makes invalidation automatic. Editing the system instructions changes the +/// tokens, the stored run stops matching, and the stale file is simply ignored. +struct PrefixCache { + + /// The exact tokens this state was produced by. + let tokens: [Int32] + let state: Edge0Model35B.State + + /// Whether `ids` begins with this cache's tokens, and so can skip them. + func matches(_ ids: [Int32]) -> Bool { + // An empty run is rejected explicitly. `prefix(0)` equals `[]` for every input, + // so without this a cache holding nothing would claim to match everything — + // currently unreachable, since neither capture nor `read` can produce one, but + // it is the kind of vacuous truth that stops being unreachable quietly. + guard !tokens.isEmpty else { return false } + + // A prompt *equal* to the prefix is rejected: replaying the last token is what + // produces the logits generation starts from, so there has to be one left. + return ids.count > tokens.count && Array(ids.prefix(tokens.count)) == tokens + } + + // MARK: - Persistence + + /// Where the snapshot lives between launches. + /// + /// **Persisting matters more than caching in memory.** An in-process cache only helps + /// the second message onward; the first one after launch — the one that decides + /// whether the app feels usable at all — still pays in full. + /// + /// **Deliberately not beside the weights.** `Documents/repacked` is created by + /// `devicectl` when the 19 GB is pushed, and the app cannot write into it: + /// + /// free=6063MiB dirExists=true + /// smallWrite=Code=513 "You don't have permission to save the file" + /// + /// Six gigabytes free and a directory that plainly exists, so this is ownership, not + /// space. The same cause explains a puzzle from the transfer work — partial files that + /// **neither** `devicectl` **nor** the app could delete. Anything the app needs to + /// write belongs somewhere the app made. + /// + /// Application Support rather than Caches: this is derived and rebuildable, but + /// rebuilding costs twenty-five seconds of the user's attention, which is not what + /// Caches is for. + /// - Parameter fingerprint: identifies the model configuration. **Required, because + /// `matches` compares tokens and nothing else.** A snapshot taken by the full + /// forty-layer model is meaningless to a pruned one — same prompt, different + /// function — and restoring it would produce confident nonsense with no error + /// anywhere. Separate files mean a configuration change simply misses the cache. + static func fileURL(in directory: URL, fingerprint: String) -> URL { + let manager = FileManager.default + guard let support = manager.urls( + for: .applicationSupportDirectory, in: .userDomainMask).first + else { + return directory.appendingPathComponent( + "prefix-cache-\(fingerprint).safetensors") + } + try? manager.createDirectory(at: support, withIntermediateDirectories: true) + return support.appendingPathComponent("prefix-cache-\(fingerprint).safetensors") + } + + /// Flattened for safetensors, which stores a flat name→array map. + /// + /// The token run is written as an array too, so the file is self-describing: nothing + /// outside it needs to stay in sync for the verification above to work. + private func flattened() -> [String: MLXArray] { + // **Everything is forced contiguous before it is written.** safetensors stores a + // plain row-major buffer and cannot represent a strided view, so saving one + // fails. The convolution state is exactly that — `state.convolution[layer]` is a + // tail slice of a larger array — and the failure was silent: the write was + // wrapped in `try?`, so the only symptom was a cache that worked within a session + // and never survived a launch. + var out: [String: MLXArray] = ["tokens": MLXArray(tokens)] + for (layer, value) in state.convolution { out["conv.\(layer)"] = contiguous(value) } + for (layer, value) in state.recurrent { out["rec.\(layer)"] = contiguous(value) } + for (layer, value) in state.keys { out["k.\(layer)"] = contiguous(value) } + for (layer, value) in state.values { out["v.\(layer)"] = contiguous(value) } + return out + } + + /// Write the snapshot. Cache persistence is best-effort; a missing cache affects + /// startup speed but not generation correctness. + func write(to url: URL) { + do { + let arrays = flattened() + eval(Array(arrays.values)) + try save(arrays: arrays, url: url) + } catch { + let directory = url.deletingLastPathComponent() + let free = (try? directory.resourceValues( + forKeys: [.volumeAvailableCapacityForImportantUsageKey]))? + .volumeAvailableCapacityForImportantUsage + let probe = directory.appendingPathComponent(".write-probe") + var probeResult = "ok" + do { + try Data([0]).write(to: probe) + try? FileManager.default.removeItem(at: probe) + } catch { + probeResult = "\(error)" + } + print(""" + [prefix-cache] write failed: \(error) + [prefix-cache] free=\(free.map { "\($0 / 1_048_576)MiB" } ?? "unknown") \ + dirExists=\(FileManager.default.fileExists(atPath: directory.path)) \ + smallWrite=\(probeResult) + """) + } + } + + /// Read a snapshot back, or `nil` if there is nothing usable there. + /// + /// Every failure mode — absent file, truncated write, a format from an older build — + /// resolves to `nil` and a normal prefill. A cache is an optimisation; it is never + /// worth failing a request over. + static func read(from url: URL) -> PrefixCache? { + guard FileManager.default.fileExists(atPath: url.path), + let arrays = try? loadArrays(url: url), + let tokenArray = arrays["tokens"] + else { return nil } + + let state = Edge0Model35B.State() + for (name, value) in arrays where name != "tokens" { + let parts = name.split(separator: ".") + guard parts.count == 2, let layer = Int(parts[1]) else { continue } + switch parts[0] { + case "conv": state.convolution[layer] = value + case "rec": state.recurrent[layer] = value + case "k": state.keys[layer] = value + case "v": state.values[layer] = value + default: break + } + } + let tokens = tokenArray.asArray(Int32.self) + guard !tokens.isEmpty else { return nil } + state.offset = tokens.count + return PrefixCache(tokens: tokens, state: state) + } +} + +extension Edge0Model35B.State { + + /// A copy that generation can advance without disturbing the original. + /// + /// The four tables are Swift dictionaries — value types — so copying them is a real + /// snapshot, and the `MLXArray`s they hold are safe to share because **every write in + /// the model reassigns rather than mutating** (`state.recurrent[layer] = ...`, all + /// four verified at their call sites). If that ever changes to an in-place update, + /// this becomes silent corruption rather than a compile error, which is why it is + /// written down here. + func snapshot() -> Edge0Model35B.State { + let copy = Edge0Model35B.State() + copy.convolution = convolution + copy.recurrent = recurrent + copy.keys = keys + copy.values = values + copy.offset = offset + return copy + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift b/ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift new file mode 100644 index 0000000..388847c --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift @@ -0,0 +1,220 @@ +import Foundation +import MLX + +/// Trained routing heads that predict expert selections for upcoming decode steps. +/// +/// During decode, these heads replace the router for their configured consumer layers; +/// the model's real gate is not evaluated for those layers. The heads and their LoRA +/// adapter are a trained pair and must be used together. Layers without a prediction use +/// the normal router. +/// +/// A head on layer N predicts routing for layer N+1 on the next token. This layer and +/// token offset allows expert reads to be prefetched before the corresponding forward +/// pass. Features contain the current hidden state and one-hot encodings of the current +/// and previous executed routes. +final class PregateHeads: @unchecked Sendable { + + enum Failure: LocalizedError { + case tensorMissing(String) + case wrongShape(String, expected: [Int], found: [Int]) + + var errorDescription: String? { + switch self { + case .tensorMissing(let name): return "pregate file has no \(name)" + case .wrongShape(let name, let expected, let found): + return "pregate \(name) is \(found), expected \(expected)" + } + } + } + + /// Layers whose head runs. Each predicts for `owner + 1`. + let owners: [Int] + /// Which layers consume a prediction, in the order their heads appear. + let consumers: [Int] + let hidden: Int + let experts: Int + + /// Keep the file mapped and materialize arrays only during a pregate forward pass. + private let shard: SafetensorsShard + private let middle: Int + + /// Where the file lives, beside the weights. Optional by construction: a build + /// without it, the normal router is used. + static func fileURL(in directory: URL, round6: Bool = false) -> URL { + if round6 { + return directory.appendingPathComponent("pregate-round6.safetensors") + } + let official = directory.appendingPathComponent("prerouter-stacked.safetensors") + if FileManager.default.fileExists(atPath: official.path) { return official } + return directory.appendingPathComponent("pregate-round6.safetensors") + } + + init(url: URL, hidden: Int, experts: Int) throws { + let shard = try SafetensorsShard(url: url) + self.shard = shard + self.hidden = hidden + self.experts = experts + let inputWidth = hidden + 2 * experts + + // `nil` in the expected shape means "whatever the file says" — the hidden width + // of the heads is the file's to decide, and everything else is pinned by the + // model it has to plug into. + func validate(_ name: String, expecting: [Int?]) throws -> SafetensorsShard.Entry { + guard let entry = shard.entries[name] else { + throw Failure.tensorMissing(name) + } + let matches = entry.shape.count == expecting.count + && zip(entry.shape, expecting).allSatisfy { $1 == nil || $0 == $1 } + guard matches else { + throw Failure.wrongShape( + name, expected: expecting.map { $0 ?? -1 }, found: entry.shape) + } + guard entry.dtype == "F16" else { + throw Failure.wrongShape(name, expected: expecting.map { $0 ?? -1 }, + found: entry.shape) + } + return entry + } + + guard let ownersEntry = shard.entries["pregate.owners"], + let ownersBase = shard.pointer(to: "pregate.owners") + else { throw Failure.tensorMissing("pregate.owners") } + let count = ownersEntry.shape.first ?? 0 + owners = (0 ..< count).map { + Int(ownersBase.load(fromByteOffset: $0 * 4, as: Int32.self)) + } + consumers = owners.map { $0 + 1 } + + let fc1 = try validate("pregate.fc1", expecting: [count, inputWidth, nil]) + middle = fc1.shape[2] + _ = try validate("pregate.fc2", expecting: [count, middle, experts]) + _ = try validate("pregate.linear_init", expecting: [count, inputWidth, experts]) + } + + private func array(_ name: String) -> MLXArray { + guard let entry = shard.entries[name], let base = shard.pointer(to: name) else { + preconditionFailure("validated pregate tensor disappeared: \(name)") + } + return MLXArray( + UnsafeRawBufferPointer(start: base, count: entry.byteCount), + entry.shape, type: Float16.self) + } + + /// One-hot rows for every owner layer at once, built on the CPU. + /// + /// **Deliberately not `identity[indices].sum(axis: 0)` per layer.** That is the + /// obvious phrasing and it was measured: thirty-three lazy MLX nodes, each pinning + /// its inputs until the step boundary evaluated them, contributed to a footprint + /// that rose from ~600 MiB to 1535 MiB. On this device that is not merely untidy — + /// the expert blocks live in *clean file-backed* pages, so anything the process + /// holds dirty is taken out of the page cache those reads depend on. `read` doubled + /// to 134 ms and the weight copy rate halved, and together they ate a 70 ms saving. + /// + /// The routing is already on the CPU — `moe` has just called `asArray` on it — so + /// the whole `[n, experts]` block is a memset and `n * k` stores, then one array. + /// No graph, nothing retained, nothing to evaluate. + func oneHotRows(_ perLayer: [[Int32]]) -> MLXArray { + var flat = [Float16](repeating: 0, count: perLayer.count * experts) + for (row, chosen) in perLayer.enumerated() { + for expert in chosen where expert >= 0 && Int(expert) < experts { + // Accumulate so duplicate IDs retain their multiplicity. + flat[row * experts + Int(expert)] += 1 + } + } + return MLXArray(flat, [perLayer.count, experts]) + } + + /// Run every head as one batch. + /// + /// - Parameters: + /// - inputs: `[n, hidden]`, the owner layers' MoE inputs this token, in owner order. + /// - executed: `[n, experts]`, the one-hot rows those layers actually ran, from + /// `oneHotRows`. + /// - previous: `[n, experts]`, the previous token's routing or a zero row. + /// - Returns: `[n, experts]` float32 logits for the consuming layers. + func logits(inputs: MLXArray, executed: MLXArray, previous: MLXArray) -> MLXArray { + let fc1 = array("pregate.fc1") + let fc2 = array("pregate.fc2") + let linearInit = array("pregate.linear_init") + let features = concatenated( + [inputs.asType(.float16), executed, previous], axis: -1) + .expandedDimensions(axis: 1) // [n, 1, in] + let inner = matmul(features, fc1) // [n, 1, mid] + // Exact erf GELU, as trained (torch's default). `geluApproximate` would be a + // different function evaluated on 33 × 512 values a step, and there is no + // reference to check it against; the tanh form is not what these weights saw. + let activated = 0.5 * inner * (1 + erf(inner / Float(2).squareRoot())) + let out = matmul(features, linearInit) + matmul(activated, fc2) + return out.squeezed(axis: 1).asType(.float32) // [n, experts] + } + + /// What has to survive from one decode step to the next. + /// + /// Per-conversation routing history used by the trained heads. + final class State { + /// Owner layer to its MoE input for the current token. Stored as an MLX array so + /// conversion to CPU does not introduce an additional synchronization point. + var moeInput: [Int: MLXArray] = [:] + /// Owner layer → the experts it actually ran this token. + var executed: [Int: [Int32]] = [:] + /// The previous token's executed experts, or absent on the first token. + var previous: [Int: [Int32]] = [:] + /// Consuming layer → the experts it will use next step, and their weights. + /// + /// Both plain Swift. They cross a step boundary, and an `MLXArray` that does + /// that carries whatever graph produced it into the next step. + var predicted: [Int: [Int]] = [:] + var scores: [Int: [Float]] = [:] + + /// Advance routing history at the end of a decode step. + func advance() { + previous = executed + executed = [:] + moeInput = [:] + } + + func reset() { + moeInput = [:]; executed = [:]; previous = [:] + predicted = [:]; scores = [:] + } + } + + // MARK: - Optional numerical self-check + + /// Compare the heads' output against an optional fixture containing inputs, expected + /// logits, and top-2 expert IDs. Returns `nil` when the fixture is unavailable. + func selfCheck(fixture url: URL) -> String? { + guard let shard = try? SafetensorsShard(url: url), + let feats = shard.entries["feats"], let featsBase = shard.pointer(to: "feats"), + let ref = shard.entries["logits"], let refBase = shard.pointer(to: "logits"), + let top = shard.entries["top2"], let topBase = shard.pointer(to: "top2") + else { return nil } + + let features = MLXArray( + UnsafeRawBufferPointer(start: featsBase, count: feats.byteCount), + feats.shape, type: Float16.self) + let expected = MLXArray( + UnsafeRawBufferPointer(start: refBase, count: ref.byteCount), + ref.shape, type: Float32.self) + let expectedTop = MLXArray( + UnsafeRawBufferPointer(start: topBase, count: top.byteCount), + top.shape, type: Int32.self) + + // The fixture holds the concatenated 2560-wide features; split them back so the + // public entry point is what gets exercised, not a shortcut around it. + let n = feats.shape[0] + let inputs = features[0..., 0 ..< hidden] + let executed = features[0..., hidden ..< (hidden + experts)] + let previous = features[0..., (hidden + experts)...] + let got = logits(inputs: inputs, executed: executed, previous: previous) + + let difference = abs(got - expected) + let gotTop = argSort(-got, axis: -1)[0..., 0 ..< 2] + let agree = (gotTop .== expectedTop).all(axis: -1).sum() + eval(difference, agree) + return String( + format: "[pregate] self-check: max|Δ|=%.4f mean|Δ|=%.5f top-2 agree %d/%d", + difference.max().item(Float.self), difference.mean().item(Float.self), + agree.item(Int32.self), Int32(n)) + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift b/ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift new file mode 100644 index 0000000..a42de63 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift @@ -0,0 +1,162 @@ +import Foundation + +/// The repacked layout: one file per layer, one expert per contiguous block. +/// +/// Each block holds an expert's nine tensors back to back, 1,769,472 bytes — exactly +/// 108 pages of 16 KB — so with data starting at offset 0 of its layer's file, **every +/// block is page-aligned**. Reading one expert is one sequential 1.69 MB read instead +/// of the nine scattered reads the original safetensors layout forces. +/// +/// Each file is checked against the expected length before it is opened so truncated +/// or incomplete files cannot be used as model weights. +final class RepackedExperts { + + /// Byte layout inside one expert's block, in the order the packer wrote them. + enum Part: Int, CaseIterable { + case gateWeight, gateScales, gateBiases + case upWeight, upScales, upBiases + case downWeight, downScales, downBiases + + var byteCount: Int { + switch self { + case .gateWeight, .upWeight, .downWeight: return 524_288 + default: return 32_768 + } + } + + var offset: Int { + Part.allCases.prefix(rawValue).reduce(0) { $0 + $1.byteCount } + } + } + + static let blockBytes = 1_769_472 + static let expertsPerLayer = 256 + static let layerBytes = blockBytes * expertsPerLayer // 452,984,832 + static let pageSize = 16384 // blockBytes is 108 of these + + private struct Opened { + let size: Int + let descriptor: Int32 + } + + // Keep descriptors, not 40 whole-file mappings. Mapping every 432 MiB layer + // consumes ~17 GiB of virtual address space and fails after 11 layers when a + // development profile lacks Extended Virtual Addressing. The hot path already + // uses pread, so those mappings were only an optional prefetch mechanism. + private var layers: [Int: Opened] = [:] + + /// Layers whose file was present but the wrong length, with the length found. + private(set) var rejected: [(layer: Int, bytes: Int)] = [] + private(set) var missing: [Int] = [] + + var availableLayers: [Int] { layers.keys.sorted() } + + init(directory: URL, expectedLayers: Int) { + for layer in 0 ..< expectedLayers { + let url = directory.appendingPathComponent(String(format: "experts-L%02d.bin", layer)) + + guard let attributes = try? FileManager.default.attributesOfItem(atPath: url.path), + let size = (attributes[.size] as? NSNumber)?.intValue + else { + missing.append(layer) + continue + } + guard size == Self.layerBytes else { + rejected.append((layer, size)) + continue + } + + let descriptor = open(url.path, O_RDONLY) + guard descriptor >= 0 else { + missing.append(layer) + continue + } + layers[layer] = Opened(size: size, descriptor: descriptor) + } + } + + deinit { + for opened in layers.values { + close(opened.descriptor) + } + } + + /// Ask the kernel for an expert's whole block in one go. + func prefetch(layer: Int, expert: Int) { + guard let opened = layers[layer], expert >= 0, expert < Self.expertsPerLayer + else { return } + + // `F_RDADVISE` is the descriptor equivalent of the old mmap + MADV_WILLNEED + // path: it starts asynchronous readahead without reserving 17 GiB of virtual + // address space. The eventual `pread` remains the source of truth, so ignored + // or failed advice can affect only latency, never model output. + var advice = radvisory( + ra_offset: off_t(expert * Self.blockBytes), + ra_count: Int32(Self.blockBytes)) + _ = withUnsafeMutablePointer(to: &advice) { + fcntl(opened.descriptor, F_RDADVISE, UnsafeMutableRawPointer($0)) + } + } + + /// Remove files that are present but have the wrong length, returning any failures. + static func removeTruncated(in directory: URL, expectedLayers: Int) + -> (removed: [(Int, Int)], failed: [(Int, String)]) + { + var removed: [(Int, Int)] = [] + var failed: [(Int, String)] = [] + for layer in 0 ..< expectedLayers { + let url = directory.appendingPathComponent(String(format: "experts-L%02d.bin", layer)) + guard let attributes = try? FileManager.default.attributesOfItem(atPath: url.path), + let size = (attributes[.size] as? NSNumber)?.intValue, + size != layerBytes + else { continue } + do { + try FileManager.default.removeItem(at: url) + removed.append((layer, size)) + } catch { + failed.append((layer, error.localizedDescription)) + } + } + return (removed, failed) + } + + /// Read one expert's whole block with a single `pread`. + /// + /// The mmap path reaches the same bytes through page faults, which the kernel + /// services in its own units; this asks for all 1,769,472 bytes in one call and + /// lets it issue one large sequential read. Worth measuring against the fault path + /// now that a block is contiguous — 768 MB/s is still only about a third of what + /// the device can do. + /// One sequential read of the whole layer file (432 MiB, 256 expert blocks). + func readLayer(_ layer: Int, into buffer: UnsafeMutableRawPointer) -> Bool { + guard let opened = layers[layer] else { return false } + var moved = 0 + while moved < Self.layerBytes { + let got = pread( + opened.descriptor, + buffer.advanced(by: moved), + Self.layerBytes - moved, + off_t(moved)) + if got <= 0 { return false } + moved += got + } + return true + } + + func read(layer: Int, expert: Int, into buffer: UnsafeMutableRawPointer) -> Bool { + guard let opened = layers[layer], expert >= 0, expert < Self.expertsPerLayer + else { return false } + + var moved = 0 + while moved < Self.blockBytes { + let got = pread( + opened.descriptor, + buffer.advanced(by: moved), + Self.blockBytes - moved, + off_t(expert * Self.blockBytes + moved)) + if got <= 0 { return false } + moved += got + } + return true + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/Safetensors.swift b/ios/Sources/Edge0MLX/Edge035B/Safetensors.swift new file mode 100644 index 0000000..c48b554 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/Safetensors.swift @@ -0,0 +1,264 @@ +import Foundation + +/// Read-only mmap over one safetensors shard, plus the header that says where each +/// tensor's bytes begin. +/// +/// A shard is: 8 bytes of little-endian header length, that many bytes of JSON, then +/// the payload. Every tensor's `data_offsets` are relative to the start of the +/// payload, so a tensor — or, for a stacked expert tensor, one expert's slice of it +/// — is a byte range and nothing more. That is the premise the whole streaming +/// design rests on: **no parsing, no dequantisation, no copy on the read path.** +/// +/// Mapped `MAP_SHARED` + read-only so the pages stay clean and file-backed. Dirtying +/// them would turn the 19 GB into the app's own memory and defeat the exercise. +final class SafetensorsShard { + + struct Entry { + let dtype: String + let shape: [Int] + /// Absolute offset in the file, payload base already added. + let offset: Int + let byteCount: Int + } + + let url: URL + private let descriptor: Int32 + private let base: UnsafeRawPointer + private let mappedSize: Int + private(set) var entries: [String: Entry] = [:] + + enum Failure: LocalizedError { + case cannotOpen(String) + case cannotMap(String) + case badHeader(String) + case truncated(String, needs: Int, has: Int) + + var errorDescription: String? { + switch self { + case .cannotOpen(let p): return "cannot open \(p)" + case .cannotMap(let p): return "cannot mmap \(p)" + case .badHeader(let p): return "bad safetensors header in \(p)" + case .truncated(let p, let needs, let has): + return "\(p) is truncated: header describes \(needs) bytes, file has \(has)" + } + } + } + + init(url: URL) throws { + self.url = url + + descriptor = open(url.path, O_RDONLY) + guard descriptor >= 0 else { throw Failure.cannotOpen(url.lastPathComponent) } + + var status = stat() + guard fstat(descriptor, &status) == 0, status.st_size > 16 else { + close(descriptor) + throw Failure.cannotOpen(url.lastPathComponent) + } + mappedSize = Int(status.st_size) + + guard let mapped = mmap(nil, mappedSize, PROT_READ, MAP_SHARED, descriptor, 0), + mapped != MAP_FAILED + else { + close(descriptor) + throw Failure.cannotMap(url.lastPathComponent) + } + base = UnsafeRawPointer(mapped) + + // Read advice is applied at the call site because access patterns vary by tensor. + + let headerLength = Int(base.loadUnaligned(as: UInt64.self)) + guard headerLength > 0, headerLength + 8 <= mappedSize else { + munmap(mapped, mappedSize) + close(descriptor) + throw Failure.badHeader(url.lastPathComponent) + } + + let payloadBase = 8 + headerLength + let headerData = Data(bytes: base.advanced(by: 8), count: headerLength) + + guard let json = try JSONSerialization.jsonObject(with: headerData) as? [String: Any] + else { + munmap(mapped, mappedSize) + close(descriptor) + throw Failure.badHeader(url.lastPathComponent) + } + + for (name, value) in json { + guard name != "__metadata__", + let fields = value as? [String: Any], + let dtype = fields["dtype"] as? String, + let shape = fields["shape"] as? [Int], + let offsets = fields["data_offsets"] as? [Int], + offsets.count == 2 + else { continue } + + entries[name] = Entry( + dtype: dtype, + shape: shape, + offset: payloadBase + offsets[0], + byteCount: offsets[1] - offsets[0] + ) + } + + // Does the header describe more bytes than the file actually holds? + // + // The device-side size check in the transfer script compares against what + // `devicectl` reports, which is three significant figures — enough to catch a + // file that arrived at 133 MB of 432, useless for catching one that is short by + // a few megabytes. This catches that case, cheaply, at the only point where it + // can still be reported as a load failure rather than as quiet nonsense several + // layers into inference. + if let end = entries.values.map({ $0.offset + $0.byteCount }).max(), end > mappedSize { + munmap(mapped, mappedSize) + close(descriptor) + throw Failure.truncated(url.lastPathComponent, needs: end, has: mappedSize) + } + } + + deinit { + munmap(UnsafeMutableRawPointer(mutating: base), mappedSize) + close(descriptor) + } + + /// Raw pointer to a whole tensor. Does not touch the pages. + func pointer(to name: String) -> UnsafeRawPointer? { + guard let entry = entries[name] else { return nil } + return base.advanced(by: entry.offset) + } + + /// Raw pointer to one slice along the leading axis of a stacked tensor. + /// + /// Expert weights are stored as `[num_experts, out, in]`, so expert *e* is one + /// contiguous run — `rowBytes` apart from its neighbour. This single line is why + /// streaming is possible at all: pulling one expert is an offset, not a load. + func pointer(to name: String, slice index: Int, of leadingDimension: Int) -> UnsafeRawPointer? { + guard let entry = entries[name], leadingDimension > 0 else { return nil } + let rowBytes = entry.byteCount / leadingDimension + guard index >= 0, index < leadingDimension else { return nil } + return base.advanced(by: entry.offset + index * rowBytes) + } + + func sliceByteCount(of name: String, leadingDimension: Int) -> Int? { + guard let entry = entries[name], leadingDimension > 0 else { return nil } + return entry.byteCount / leadingDimension + } + + /// Reset the kernel's read-ahead policy for the whole mapping. + func advise(_ advice: Int32) { + madvise(UnsafeMutableRawPointer(mutating: base), mappedSize, advice) + } + + /// Ask the kernel to start pulling one expert's slice in. + /// + /// `MADV_WILLNEED` is advice, not a read: it returns once the I/O is queued rather + /// than once the pages are there. Issued immediately before the copy it can only + /// overlap within the slice; issued a step ahead — which is what pregate makes + /// possible — it has a whole layer of compute to land in. + func prefetch(_ name: String, slice index: Int, of leadingDimension: Int) { + guard let entry = entries[name], leadingDimension > 0, + index >= 0, index < leadingDimension + else { return } + let rowBytes = entry.byteCount / leadingDimension + madvise( + UnsafeMutableRawPointer(mutating: base).advanced(by: entry.offset + index * rowBytes), + rowBytes, MADV_WILLNEED) + } + + /// Ask the kernel to bring one whole tensor in. + /// + /// Request that the kernel warm the tensor's clean, file-backed pages. + /// Returns the number of bytes requested. + @discardableResult + func warm(_ name: String) -> Int { + guard let entry = entries[name] else { return 0 } + madvise(UnsafeMutableRawPointer(mutating: base).advanced(by: entry.offset), + entry.byteCount, MADV_WILLNEED) + return entry.byteCount + } +} + +/// The four shards plus the index that says which shard holds which tensor. +final class SafetensorsBundle { + + private var shards: [String: SafetensorsShard] = [:] + private var shardForTensor: [String: String] = [:] + + let directory: URL + + init(directory: URL) throws { + self.directory = directory + + let indexURL = directory.appendingPathComponent("model.safetensors.index.json") + let indexData = try Data(contentsOf: indexURL) + guard let root = try JSONSerialization.jsonObject(with: indexData) as? [String: Any], + let map = root["weight_map"] as? [String: String] + else { + throw SafetensorsShard.Failure.badHeader("model.safetensors.index.json") + } + shardForTensor = map + + for file in Set(map.values) { + shards[file] = try SafetensorsShard(url: directory.appendingPathComponent(file)) + } + } + + var mappedShardCount: Int { shards.count } + + /// Every tensor name in the index, and which shard holds it. + var index: [String: String] { shardForTensor } + + func entry(_ name: String) -> SafetensorsShard.Entry? { + guard let file = shardForTensor[name] else { return nil } + return shards[file]?.entries[name] + } + + func pointer(to name: String) -> UnsafeRawPointer? { + guard let file = shardForTensor[name] else { return nil } + return shards[file]?.pointer(to: name) + } + + func pointer(to name: String, slice index: Int, of leadingDimension: Int) -> UnsafeRawPointer? { + guard let file = shardForTensor[name] else { return nil } + return shards[file]?.pointer(to: name, slice: index, of: leadingDimension) + } + + func sliceByteCount(of name: String, leadingDimension: Int) -> Int? { + guard let file = shardForTensor[name] else { return nil } + return shards[file]?.sliceByteCount(of: name, leadingDimension: leadingDimension) + } + + func advise(_ advice: Int32) { + for shard in shards.values { shard.advise(advice) } + } + + /// Bring in every tensor the decode path reads on **every** step. + /// + /// Deliberately not "the whole mapping": these shards hold the 256 routed experts per + /// layer as well, nineteen gigabytes of data that only two-of-256 of is touched per + /// layer per token. Asking for that would evict everything worth having. + /// + /// Two exclusions, both by measurement rather than taste: + /// + /// - `switch_mlp` — the routed experts, streamed on demand and prefetched from their + /// own repacked files by their actual routing. + /// - `embed_tokens` — 273 MiB of table from which `rows(_:_:)` copies one row per + /// token. Warming it would be 273 MiB of page cache for 1,152 bytes of need. + /// + /// What remains is what `weight(_:)` copies every step: 20.4 MiB per layer of + /// attention and shared-expert projections, plus `lm_head`. + @discardableResult + func warmResidentWeights() -> Int { + var total = 0 + for (name, file) in shardForTensor + where !name.contains("switch_mlp") && !name.contains("embed_tokens") { + total += shards[file]?.warm(name) ?? 0 + } + return total + } + + func prefetch(_ name: String, slice index: Int, of leadingDimension: Int) { + guard let file = shardForTensor[name] else { return } + shards[file]?.prefetch(name, slice: index, of: leadingDimension) + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/StepProfile.swift b/ios/Sources/Edge0MLX/Edge035B/StepProfile.swift new file mode 100644 index 0000000..19e8fa6 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/StepProfile.swift @@ -0,0 +1,148 @@ +import Foundation + +/// Tracks synchronous portions of model steps without forcing additional MLX evaluation. +/// GPU operations remain lazy, so timing them separately would change execution order. +final class StepProfile { + + private(set) var syncLinear: Double = 0 + private(set) var syncFull: Double = 0 + var routerSync: Double { syncLinear + syncFull } + private(set) var expertRead: Double = 0 + private(set) var expertStack: Double = 0 + + /// Time spent constructing MLX arrays from mapped weight bytes. + private(set) var weightCopy: Double = 0 + private(set) var weightBytes: Int = 0 + private(set) var expertBytes: Int = 0 + + /// Number of expert read batches and individual reads issued. + private(set) var readBatches = 0 + private(set) var readsIssued = 0 + + /// Number of unique experts requested compared with the naive T*K upper bound. + private(set) var unionTotal = 0 + private(set) var pairsTotal = 0 + + /// How often the "same experts as last token" guess was right. + /// + /// The whole prefetch scheme rests on this. It costs nothing when wrong, so any + /// positive rate is upside — but it bounds the gain: the read can only be hidden for + /// the fraction of blocks the guess anticipated. + private(set) var prefetchHits = 0 + private(set) var prefetchTotal = 0 + + /// Cumulative cost of the batched head forward plus its single barrier. + private(set) var pregateStage = 0.0 + /// Core ML vocabulary projection plus its 4 KiB MLX-to-Core ML handoff. + private(set) var vocabularyHead = 0.0 + + private var steps = 0 + + func reset() { + syncLinear = 0; syncFull = 0; expertRead = 0; expertStack = 0; expertBytes = 0 + weightCopy = 0; weightBytes = 0; pregateStage = 0; vocabularyHead = 0 + prefetchHits = 0; prefetchTotal = 0 + readBatches = 0; readsIssued = 0; steps = 0 + unionTotal = 0; pairsTotal = 0 + } + + func countStep() { steps += 1 } + + /// Timing categories for synchronous portions of a model step. + enum Phase { + case syncLinear, syncFull, expertRead, expertStack, weightCopy + /// Synchronous work performed by the trained routing heads. + case pregateStage, vocabularyHead + } + + /// Time a synchronous block. Returns the block's value so call sites stay expressions. + @inline(__always) + func measure(_ phase: Phase, _ body: () throws -> T) rethrows -> T { + let start = DispatchTime.now().uptimeNanoseconds + let value = try body() + let elapsed = Double(DispatchTime.now().uptimeNanoseconds - start) / 1_000_000 + switch phase { + case .syncLinear: syncLinear += elapsed + case .syncFull: syncFull += elapsed + case .expertRead: expertRead += elapsed + case .expertStack: expertStack += elapsed + case .weightCopy: weightCopy += elapsed + case .pregateStage: pregateStage += elapsed + case .vocabularyHead: vocabularyHead += elapsed + } + return value + } + + @inline(__always) + func measure(_ phase: Phase, _ body: () async throws -> T) async rethrows -> T { + let start = DispatchTime.now().uptimeNanoseconds + let value = try await body() + let elapsed = Double(DispatchTime.now().uptimeNanoseconds - start) / 1_000_000 + switch phase { + case .syncLinear: syncLinear += elapsed + case .syncFull: syncFull += elapsed + case .expertRead: expertRead += elapsed + case .expertStack: expertStack += elapsed + case .weightCopy: weightCopy += elapsed + case .pregateStage: pregateStage += elapsed + case .vocabularyHead: vocabularyHead += elapsed + } + return value + } + + func recordRouting(union: Int, pairs: Int) { + unionTotal += union + pairsTotal += pairs + } + + func recordPrefetch(hits: Int, total: Int) { + prefetchHits += hits + prefetchTotal += total + } + + func recordWeight(bytes: Int) { weightBytes += bytes } + + func recordRead(bytes: Int, issued: Int) { + expertBytes += bytes + readBatches += 1 + readsIssued += issued + } + + /// Per-step medians are not available without keeping every sample; per-step means + /// are, and the question here is where a quarter-second goes, not its distribution. + var report: String { + guard steps > 0 else { return "no steps" } + let n = Double(steps) + func ms(_ total: Double) -> String { String(format: "%.0f", total / n) } + let throughput = expertRead > 0 + ? String(format: "%.0f", Double(expertBytes) / (expertRead / 1000) / 1_000_000) + : "-" + let depth = readBatches > 0 + ? String(format: "%.2f", Double(readsIssued) / Double(readBatches)) + : "-" + let copyRate = weightCopy > 0 + ? String(format: "%.0f", Double(weightBytes) / (weightCopy / 1000) / 1_000_000) + : "-" + let hitRate = prefetchTotal > 0 + ? String(format: "%.0f%%", + Double(prefetchHits) / Double(prefetchTotal) * 100) + : "-" + let sharing = pairsTotal > 0 + ? String(format: "%.2f", Double(unionTotal) / Double(pairsTotal)) + : "-" + return "syncLinear=\(ms(syncLinear))ms syncFull=\(ms(syncFull))ms " + + "read=\(ms(expertRead))ms " + + "stack=\(ms(expertStack))ms readMBps=\(throughput) depth=\(depth) " + + "union/pairs=\(sharing) " + + "weightCopy=\(ms(weightCopy))ms wMBps=\(copyRate)" + // Only when the heads are installed, so a baseline build's line is + // byte-identical to the one every earlier measurement was read from. + + (pregateStage > 0 ? " pregate=\(ms(pregateStage))ms" : "") + + (vocabularyHead > 0 ? " coremlHead=\(ms(vocabularyHead))ms" : "") + // `hitRate` was computed and dropped on the floor by an earlier version of + // this line. It is the number that bounds the whole prefetch scheme — the + // read can only be hidden for the fraction of blocks the guess anticipated + // — so leaving it unprinted made the one decision it informs unanswerable. + + (prefetchTotal > 0 ? " prefetch=\(hitRate)" : "") + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift b/ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift new file mode 100644 index 0000000..5d1f377 --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift @@ -0,0 +1,57 @@ +import Foundation + +/// Separates the model's reasoning from its reply. +/// +/// This checkpoint is opened with `<|im_start|>assistant\n\n`, so **every answer +/// begins inside a reasoning block** and the first thing it writes is working-out, not a +/// reply. The closing `` is generated; the opening tag never is, because it came +/// from the prompt. +/// +/// That shape leaks straight through a stream that treats decoded text as the answer — +/// the user sees the model's notes to itself, followed by a bare ``, followed by +/// the reply. +/// +/// **The cost of hiding it is that nothing appears until the model stops reasoning.** At +/// a few tokens per second that is a real wait with an empty bubble at the end of it; the +/// generating indicator is what carries that interval. Showing the reasoning instead was +/// the alternative, and it is worse for the ordinary case of wanting an answer. +enum ThinkingBlock { + + static let terminator = "" + + /// The part of `raw` the user should see, or `nil` while the model is still reasoning. + /// + /// Leading whitespace after the terminator is dropped: the model writes a blank line + /// or two before the reply proper, which would otherwise open every answer with a gap. + static func visible(in raw: String) -> String? { + guard let end = raw.range(of: terminator) else { return nil } + return String(raw[end.upperBound...]) + .trimmingCharacters(in: .whitespacesAndNewlines) + } + + /// The non-thinking prompt normally starts directly in the answer, but some + /// generations still emit the checkpoint's literal thinking wrapper. Hide that + /// block while it is incomplete, then expose only the text after its terminator. + static func visibleReply(in raw: String) -> String? { + guard raw.contains("") || raw.contains(terminator) else { return raw } + return visible(in: raw) + } + + /// What to show once generation has finished. + /// + /// **Falls back to the whole text when the block never closed.** A reply truncated by + /// the token limit mid-reasoning has no terminator, and showing nothing at all would + /// present a completed generation as an empty answer. Raw reasoning is poor output; + /// an empty bubble is a bug report. + static func settled(_ raw: String) -> String { + visible(in: raw) ?? raw.trimmingCharacters(in: .whitespacesAndNewlines) + } + + /// Final text for a prompt that requested a direct answer. + static func settledReply(_ raw: String) -> String { + guard raw.contains("") || raw.contains(terminator) else { + return raw.trimmingCharacters(in: .whitespacesAndNewlines) + } + return visible(in: raw) ?? "" + } +} diff --git a/ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift b/ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift new file mode 100644 index 0000000..8020dee --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift @@ -0,0 +1,267 @@ +import Foundation + +/// Byte-level BPE, the GPT-4 variant this checkpoint uses. +/// +/// The pipeline, in order, because getting the order wrong produces a tokenizer that +/// works on English and mis-splits everything else: +/// +/// 1. **NFC** normalisation +/// 2. **Split** on the pre-tokenizer regex, `Isolated` — each match becomes its own piece +/// 3. **ByteLevel** — every UTF-8 byte maps to one printable character, so BPE never sees +/// a raw byte and no `unk` token is needed +/// 4. **BPE** — repeatedly merge the adjacent pair with the lowest rank +/// +/// Loaded from the compact binary produced by `tools/convert_tokenizer.py`: 5.6 MiB of +/// flat arrays rather than 19 MiB of JSON, with merges stored as **token id pairs** so +/// the merge loop never looks anything up by string. +final class Tokenizer { + + enum Failure: LocalizedError { + case cannotOpen(String) + case badHeader(String) + case truncated(String, needs: Int, has: Int) + + var errorDescription: String? { + switch self { + case .cannotOpen(let p): return "cannot open \(p)" + case .badHeader(let p): return "bad tokenizer header in \(p)" + case .truncated(let p, let needs, let has): + return "\(p) is truncated: needs \(needs) bytes, has \(has)" + } + } + } + + private let tokens: [String] // id -> token text (byte-level encoded) + private let ids: [String: Int32] // token text -> id + private let ranks: [UInt64: Int32] // packed (left, right) -> merge rank + private let specials: [String: Int32] + private let specialTexts: [Int32: String] + private let splitter: NSRegularExpression + + /// Byte ↔ printable-character tables, the GPT-2 mapping. + /// + /// Bytes that are already printable map to themselves; the rest are shifted into an + /// unused Unicode range. This is what lets a BPE vocabulary of *characters* represent + /// arbitrary bytes without an unknown token. + private static let byteToUnicode: [Character] = { + var mapping: [Int] = [] + var table = [Character](repeating: " ", count: 256) + for value in UInt8(ascii: "!") ... UInt8(ascii: "~") { mapping.append(Int(value)) } + for value in 0xA1 ... 0xAC { mapping.append(value) } + for value in 0xAE ... 0xFF { mapping.append(value) } + + var extra = 0 + for byte in 0 ..< 256 { + if mapping.contains(byte) { + table[byte] = Character(UnicodeScalar(byte)!) + } else { + table[byte] = Character(UnicodeScalar(256 + extra)!) + extra += 1 + } + } + return table + }() + + private static let unicodeToByte: [Character: UInt8] = { + var inverse: [Character: UInt8] = [:] + for byte in 0 ..< 256 { inverse[byteToUnicode[byte]] = UInt8(byte) } + return inverse + }() + + init(url: URL) throws { + guard let data = try? Data(contentsOf: url, options: .mappedIfSafe), + data.count > 8 + else { throw Failure.cannotOpen(url.lastPathComponent) } + + let headerLength = Int(data.withUnsafeBytes { $0.loadUnaligned(as: UInt64.self) }) + guard headerLength > 0, 8 + headerLength <= data.count, + let header = try? JSONSerialization.jsonObject( + with: data[8 ..< (8 + headerLength)]) as? [String: Any], + let vocabSize = header["vocab_size"] as? Int, + let mergeCount = header["merge_count"] as? Int, + let specialCount = header["special_count"] as? Int, + let pattern = header["regex"] as? String + else { throw Failure.badHeader(url.lastPathComponent) } + + // Sections are fixed-width and laid out in a known order, so each is a slice at a + // running offset. Checked against the file length before any of it is read — a + // truncated tokenizer that silently loses its tail would encode most text + // correctly and fail only on rare tokens. + var cursor = 8 + headerLength + func take(_ count: Int) throws -> Data { + guard cursor + count <= data.count else { + throw Failure.truncated(url.lastPathComponent, + needs: cursor + count, has: data.count) + } + defer { cursor += count } + return data[cursor ..< (cursor + count)] + } + + let offsets = try take((vocabSize + 1) * 4).toUInt32Array() + let blob = try take(Int(offsets[vocabSize])) + let mergeFlat = try take(mergeCount * 8).toUInt32Array() + let specialIds = try take(specialCount * 4).toUInt32Array() + let specialOffsets = try take((specialCount + 1) * 4).toUInt32Array() + let specialBlob = try take(Int(specialOffsets[specialCount])) + + var tokens = [String](repeating: "", count: vocabSize) + var ids: [String: Int32] = [:] + ids.reserveCapacity(vocabSize) + blob.withUnsafeBytes { raw in + let base = raw.bindMemory(to: UInt8.self).baseAddress! + for index in 0 ..< vocabSize { + let start = Int(offsets[index]), end = Int(offsets[index + 1]) + let text = String(decoding: UnsafeBufferPointer( + start: base + start, count: end - start), as: UTF8.self) + tokens[index] = text + ids[text] = Int32(index) + } + } + self.tokens = tokens + self.ids = ids + + var ranks: [UInt64: Int32] = [:] + ranks.reserveCapacity(mergeCount) + for rank in 0 ..< mergeCount { + let left = UInt64(mergeFlat[rank * 2]) + let right = UInt64(mergeFlat[rank * 2 + 1]) + ranks[(left << 32) | right] = Int32(rank) + } + self.ranks = ranks + + var specials: [String: Int32] = [:] + var specialTexts: [Int32: String] = [:] + specialBlob.withUnsafeBytes { raw in + let base = raw.bindMemory(to: UInt8.self).baseAddress! + for index in 0 ..< specialCount { + let start = Int(specialOffsets[index]), end = Int(specialOffsets[index + 1]) + let text = String(decoding: UnsafeBufferPointer( + start: base + start, count: end - start), as: UTF8.self) + specials[text] = Int32(specialIds[index]) + specialTexts[Int32(specialIds[index])] = text + } + } + self.specials = specials + self.specialTexts = specialTexts + + // Carried across from the source file rather than rewritten, because this pattern + // is where CJK and emoji splitting is decided. + self.splitter = try NSRegularExpression(pattern: pattern, options: []) + } + + var vocabularySize: Int { tokens.count } + var specialTokenCount: Int { specials.count } + + func id(of special: String) -> Int32? { specials[special] } + + // MARK: - Encode + + func encode(_ text: String) -> [Int32] { + // Special tokens are matched before anything else and never split — a chat + // template's `<|im_start|>` reaching the BPE loop would come back as a handful of + // ordinary pieces that happen to spell it. + var output: [Int32] = [] + for (chunk, special) in splitOnSpecials(text) { + if let special { + output.append(special) + } else { + output.append(contentsOf: encodeOrdinary(chunk)) + } + } + return output + } + + private func splitOnSpecials(_ text: String) -> [(String, Int32?)] { + guard !specials.isEmpty else { return [(text, nil)] } + var pieces: [(String, Int32?)] = [] + var remainder = Substring(text) + + while !remainder.isEmpty { + var earliest: (Range, Int32)? = nil + for (token, id) in specials { + if let range = remainder.range(of: token), + earliest == nil || range.lowerBound < earliest!.0.lowerBound { + earliest = (range, id) + } + } + guard let (range, id) = earliest else { break } + if range.lowerBound > remainder.startIndex { + pieces.append((String(remainder[remainder.startIndex ..< range.lowerBound]), nil)) + } + pieces.append(("", id)) + remainder = remainder[range.upperBound...] + } + if !remainder.isEmpty { pieces.append((String(remainder), nil)) } + return pieces + } + + private func encodeOrdinary(_ text: String) -> [Int32] { + let normalised = text.precomposedStringWithCanonicalMapping // NFC + var output: [Int32] = [] + + let full = NSRange(normalised.startIndex ..< normalised.endIndex, in: normalised) + splitter.enumerateMatches(in: normalised, options: [], range: full) { match, _, _ in + guard let match, let range = Range(match.range, in: normalised) else { return } + let piece = String(normalised[range]) + + // Byte level: every UTF-8 byte becomes one printable character. + let mapped = String(piece.utf8.map { Self.byteToUnicode[Int($0)] }) + output.append(contentsOf: bpe(mapped)) + } + return output + } + + /// Merge the lowest-ranked adjacent pair until none remains. + private func bpe(_ text: String) -> [Int32] { + if let whole = ids[text] { return [whole] } + + var parts = text.map { String($0) }.compactMap { ids[$0] } + guard parts.count > 1 else { return parts } + + while parts.count > 1 { + var bestRank = Int32.max + var bestIndex = -1 + for index in 0 ..< (parts.count - 1) { + let key = (UInt64(UInt32(bitPattern: parts[index])) << 32) + | UInt64(UInt32(bitPattern: parts[index + 1])) + if let rank = ranks[key], rank < bestRank { + bestRank = rank + bestIndex = index + } + } + guard bestIndex >= 0 else { break } + + let merged = tokens[Int(parts[bestIndex])] + tokens[Int(parts[bestIndex + 1])] + guard let id = ids[merged] else { break } + parts.replaceSubrange(bestIndex ... (bestIndex + 1), with: [id]) + } + return parts + } + + // MARK: - Decode + + func decode(_ sequence: [Int32]) -> String { + var bytes: [UInt8] = [] + for id in sequence { + if let special = specialTexts[id] { + bytes.append(contentsOf: Array(special.utf8)) + continue + } + guard id >= 0, Int(id) < tokens.count else { continue } + for character in tokens[Int(id)] { + if let byte = Self.unicodeToByte[character] { bytes.append(byte) } + } + } + return String(decoding: bytes, as: UTF8.self) + } +} + +private extension Data { + func toUInt32Array() -> [UInt32] { + withUnsafeBytes { raw in + Array(UnsafeBufferPointer( + start: raw.bindMemory(to: UInt32.self).baseAddress!, + count: count / 4)) + } + } +} diff --git a/ios/Sources/Edge0MLX/Edge0ChatEngine.swift b/ios/Sources/Edge0MLX/Edge0ChatEngine.swift new file mode 100644 index 0000000..c5515da --- /dev/null +++ b/ios/Sources/Edge0MLX/Edge0ChatEngine.swift @@ -0,0 +1,249 @@ +import Edge0Core +import Foundation +import MLX + +public struct Edge0GenerationResult: Sendable { + public let text: String + public let generatedTokenCount: Int + public let elapsedSeconds: TimeInterval + public let prefillSeconds: TimeInterval + public let prefillTokensPerSecond: Double + public let timeToFirstTokenSeconds: TimeInterval + public let decodeTokensPerSecond: Double + public let peakMemoryBytes: Int +} + +private enum Edge0ChatTemplate8B { + static func firstTurn(_ text: String, thinking: Bool) -> String { + "SYSTEMdetailed thinking \(thinking ? "on" : "off")<|role_end|>" + + "HUMAN\(text)<|role_end|>" + assistantPrefix(thinking: thinking) + } + + static func nextTurn(_ text: String, thinking: Bool) -> String { + "<|role_end|>HUMAN\(text)<|role_end|>" + + assistantPrefix(thinking: thinking) + } + + private static func assistantPrefix(thinking: Bool) -> String { + "ASSISTANT\n" + (thinking ? "" : "") + } +} + +/// Product-facing, stateful chat entry point. The model math and cache lifetime +/// remain owned by StreamingEdge0Model8B; this type only mirrors the official +/// chat template and official sampling loop. +public final class Edge0ChatEngine: @unchecked Sendable { + private let tokenizer: Edge0Tokenizer8B + private let model: StreamingEdge0Model8B + private var hasConversationContext = false + private var tokenCount = 0 + + /// Short prompts do not amortize loading all 128 experts per layer. + private static let fullLayerPrefillMinimum = 128 + + private static var runtimeDevice: Device { + #if targetEnvironment(simulator) + .cpu + #else + .gpu + #endif + } + + public init(modelURL: URL, progress: (String) -> Void = { _ in }) throws { + let built = try Device.withDefaultDevice(Self.runtimeDevice) { + let configuration = try Edge0Configuration8B.load(from: modelURL.appendingPathComponent("config.json")) + let tokenizer = try Edge0Tokenizer8B(contentsOf: modelURL.appendingPathComponent("tokenizer.json")) + let store = try ExpertTensorStore( + modelURL: modelURL.appendingPathComponent("model.safetensors"), + expertCount: configuration.numExperts + ) + let model = try StreamingEdge0Model8B( + configuration: configuration, store: store, + loraWeightsURL: modelURL.appendingPathComponent( + "lora_edge0_8b.safetensors"), + prerouterWeightsURL: modelURL.appendingPathComponent( + "prerouter_edge0_8b.safetensors"), + progress: progress) + progress("prewarming official decode path") + try model.prewarm() + return (tokenizer, model) + } + tokenizer = built.0 + model = built.1 + } + + public func reset() { + model.reset() + hasConversationContext = false + tokenCount = 0 + } + + public func reply( + to userText: String, + maxTokens: Int = 2048, + thinking: Bool = false, + seed: UInt64? = nil, + onText: @escaping @Sendable (String) -> Void = { _ in }, + shouldContinue: @escaping @Sendable () -> Bool = { true } + ) throws -> Edge0GenerationResult { + try Device.withDefaultDevice(Self.runtimeDevice) { + try generateReply(to: userText, maxTokens: maxTokens, + thinking: thinking, seed: seed, + onText: onText, shouldContinue: shouldContinue) + } + } + + private func generateReply( + to userText: String, maxTokens: Int, thinking: Bool, seed: UInt64?, + onText: @escaping @Sendable (String) -> Void, + shouldContinue: @escaping @Sendable () -> Bool + ) throws -> Edge0GenerationResult { + let trimmed = userText.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { throw M1Error.invalid("Message cannot be empty") } + guard maxTokens > 0 else { throw M1Error.invalid("maxTokens must be positive") } + + let prompt = hasConversationContext + ? Edge0ChatTemplate8B.nextTurn(trimmed, thinking: thinking) + : Edge0ChatTemplate8B.firstTurn(trimmed, thinking: thinking) + let promptIDs = try tokenizer.encode(prompt) + guard !promptIDs.isEmpty else { throw M1Error.invalid("Tokenizer returned an empty prompt") } + + let started = Date() + guard shouldContinue() else { + model.reset(); hasConversationContext = false; tokenCount = 0 + throw CancellationError() + } + var output: StreamingModelOutput? + // Full-layer prefill loads every expert in a layer. That cost is almost + // independent of length, so a short chat prompt pays for the whole 4 GB + // expert set. Below this cutoff, walk the same top-8 decode path. + let streamingPrefill = promptIDs.count < Self.fullLayerPrefillMinimum + if streamingPrefill { + for id in promptIDs { + guard shouldContinue() else { + model.reset(); hasConversationContext = false; tokenCount = 0 + throw CancellationError() + } + output = try model(tokenID: id) + } + } else { + let prefillChunk = 2048 + for start in stride(from: 0, to: promptIDs.count, by: prefillChunk) { + guard shouldContinue() else { + model.reset(); hasConversationContext = false; tokenCount = 0 + throw CancellationError() + } + let end = min(start + prefillChunk, promptIDs.count) + let chunk = Array(promptIDs[start.. 1 { + output = try model.prefill(tokenIDs: chunk) + } else if let id = chunk.first { + output = try model(tokenID: id) + } + } + } + let prefillFinished = Date() + tokenCount += promptIDs.count + + var generated: [Int] = [] + var samplingHistory = promptIDs + var peakMemory = output?.peakMemory ?? 0 + var firstTokenAt: Date? + var decodeStartedAt: Date? + while generated.count < maxTokens, let current = output { + if !shouldContinue() { + if generated.isEmpty { + model.reset(); hasConversationContext = false; tokenCount = 0 + throw CancellationError() + } + break + } + let next = try sample(current.logits, history: samplingHistory, + firstToken: generated.isEmpty, seed: seed.map { + $0 &+ UInt64(generated.count) + }) + if next == tokenizer.endOfTurnTokenID { break } + generated.append(next) + samplingHistory.append(next) + onText(try tokenizer.decode(generated)) + if firstTokenAt == nil { + firstTokenAt = Date() + decodeStartedAt = firstTokenAt + } + if generated.count == maxTokens { break } + output = try model(tokenID: next) + peakMemory = max(peakMemory, output?.peakMemory ?? 0) + tokenCount += 1 + } + hasConversationContext = true + + let finished = Date() + let firstTokenSeconds = firstTokenAt?.timeIntervalSince(started) + ?? finished.timeIntervalSince(started) + let decodeTokenCount = max(0, generated.count - 1) + let decodeSeconds = decodeStartedAt.map { finished.timeIntervalSince($0) } ?? 0 + let prefillSeconds = prefillFinished.timeIntervalSince(started) + let decodeTokensPerSecond = decodeSeconds > 0 + ? Double(decodeTokenCount) / decodeSeconds : 0 + print(String( + format: "[8b prefill] path=%@ tokens=%d total=%.2fs rate=%.2f tok/s", + streamingPrefill ? "streaming-top8" : "full-layer", + promptIDs.count, + prefillSeconds, + Double(promptIDs.count) / max(0.000_001, prefillSeconds))) + print(String( + format: "[8b decode] tokens=%d total=%.2fs rate=%.2f tok/s peak=%.0f MB", + decodeTokenCount, + decodeSeconds, + decodeTokensPerSecond, + Double(peakMemory) / 1_048_576)) + + return Edge0GenerationResult( + text: try tokenizer.decode(generated), + generatedTokenCount: generated.count, + elapsedSeconds: finished.timeIntervalSince(started), + prefillSeconds: prefillSeconds, + prefillTokensPerSecond: Double(promptIDs.count) / + max(0.000_001, prefillSeconds), + timeToFirstTokenSeconds: firstTokenSeconds, + decodeTokensPerSecond: decodeTokensPerSecond, + peakMemoryBytes: peakMemory + ) + } + + /// Official Edge0 policy: first token greedy, then temperature 0.7, + /// top-k 64, top-p 0.95 and HF-style repetition penalty 1.1. + private func sample(_ input: MLXArray, history: [Int], + firstToken: Bool, seed: UInt64?) throws -> Int { + let logits = input.flattened().asType(.float32) + if firstToken { return argMax(logits).item(Int32.self).intValue } + + if !history.isEmpty { + let unique = Array(Set(history)).sorted() + let ids = MLXArray(unique.map(Int32.init)) + let values = logits[ids] + logits[ids] = which(values .> 0, values / 1.1, values * 1.1) + } + let scaled = logits / 0.7 + let k = min(64, scaled.dim(0)) + let candidateIDs = argPartition(-scaled, kth: k - 1)[.. (indices: MLXArray, weights: MLXArray) { + precondition(nGroup > 0) + precondition(logits.shape == expertBias.shape || expertBias.shape == [logits.dim(-1)]) + precondition(logits.dim(-1) / nGroup >= 2) + precondition(topK <= topkGroup * (logits.dim(-1) / nGroup)) + let raw = sigmoid(logits.asType(.float32)) + var selection = raw + expertBias + let experts = selection.dim(-1) + precondition(experts % nGroup == 0) + precondition(topK > 0 && topK <= experts) + precondition(topkGroup > 0 && topkGroup <= nGroup) + + let groupsToDrop = nGroup - topkGroup + if groupsToDrop > 0 { + let grouped = unflatten(selection, axis: -1, shape: [nGroup, -1]) + let groupScores = top(grouped, k: 2, axis: -1).sum(axis: -1, keepDims: true) + let dropped = argPartition( + groupScores, + kth: groupsToDrop - 1, + axis: -2 + )[.ellipsis, .. MLXArray { + try stacked(ownerIDs.map { + try store.mlxArrayCopying(named: "layers.\($0).\(suffix)") + }, axis: 0) + } + fc1 = try weights("fc1.weight") + fc2 = try weights("fc2.weight") + linear = try weights("linear_init.weight") + eval(fc1, fc2, linear) + } + + func predict(features: [Int: LingPrerouterFeature]) throws + -> [Int: PredictedMoERoute] { + var activeOwners: [Int] = [] + var inputs: [MLXArray] = [] + var currentIndices: [Int: MLXArray] = [:] + + for owner in owners { + guard let feature = features[owner] else { continue } + let rows = feature.routeIndices.dim(0) + guard rows > 0 else { continue } + let current = row(feature.routeIndices, at: rows - 1) + let previous = previousIndices[owner] + ?? (rows > 1 ? row(feature.routeIndices, at: rows - 2) : nil) + currentIndices[owner] = current + + let tokenInput = lastToken(feature.input) + let combined = concatenated([ + tokenInput, + oneHot(current), + oneHot(previous), + ], axis: -1) + activeOwners.append(owner) + inputs.append(combined) + } + guard activeOwners == owners else { return [:] } + + // Official `head_batch`: 16 owners share three batched matmuls rather + // than issuing 48 tiny per-head GPU operations. + let x = stacked(inputs, axis: 0).asType(fc1.dtype) + let hidden = matmul(x, fc1.transposed(0, 2, 1)) + let gelu = 0.5 * hidden * (1 + erf(hidden / Float(2).squareRoot())) + let logits = matmul(gelu, fc2.transposed(0, 2, 1)) + + matmul(x, linear.transposed(0, 2, 1)) + let route = Edge0GroupedRouter8B.select( + logits: logits, expertBias: MLXArray.zeros([c.numExperts]), + topK: c.expertsPerToken, nGroup: c.nGroup, + topkGroup: c.topkGroup, normalize: c.normTopkProb, + routedScale: c.routedScalingFactor) + eval(route.indices, route.weights) + let allIDs = route.indices.asArray(Int32.self).map(Int.init) + let allWeights = route.weights.asArray(Float.self) + + var predictions: [Int: PredictedMoERoute] = [:] + for (position, owner) in owners.enumerated() { + let start = position * c.expertsPerToken + let end = start + c.expertsPerToken + predictions[owner + 1] = PredictedMoERoute( + expertIndices: Array(allIDs[start.. MLXArray { + var values = MLXArray.zeros([c.numExperts]) + guard let ids else { return values.reshaped([1, c.numExperts]) } + values[ids.flattened()] = MLXArray.ones([c.expertsPerToken]) + return values.reshaped([1, c.numExperts]) + } + + private func row(_ values: MLXArray, at index: Int) -> MLXArray { + let parts = split(values, indices: [index, index + 1], axis: 0) + return parts[1] + } + + private func lastToken(_ input: MLXArray) -> MLXArray { + guard input.dim(0) > 1 else { return input } + return split(input, indices: [input.dim(0) - 1], axis: 0)[1] + } + + func reset() { previousIndices.removeAll(keepingCapacity: true) } +} diff --git a/ios/Sources/Edge0MLX/MLXExpertTensorStore.swift b/ios/Sources/Edge0MLX/MLXExpertTensorStore.swift new file mode 100644 index 0000000..3e7962a --- /dev/null +++ b/ios/Sources/Edge0MLX/MLXExpertTensorStore.swift @@ -0,0 +1,56 @@ +import Edge0Core +import Foundation +import MLX + +public enum Edge0MLXError: Error { + case unsupportedDType(SafetensorsDType) +} + +public extension ExpertTensorStore { + /// Mirrors the upstream streaming loader's targeted `MADV_WILLNEED` + /// pass. It retains no MLX arrays and is safe to issue speculatively. + func adviseExpertsWillNeed(layer: Int, experts: [Int]) { + for expert in Set(experts) where (0.. MLXArray { + let slice = try slice(layer: layer, expert: expert, projection: projection, part: part) + return try copyTensor(descriptor: slice.descriptor, pointer: slice.pointer) + } + + /// Copies one row from an arbitrary axis-0-contiguous checkpoint tensor. + /// Used by the quantized embedding so a token lookup never materializes + /// the full 157k x 1536 embedding table. + func mlxAxisZeroSliceCopying(named name: String, index row: Int) throws -> MLXArray { + let tensor = try index.tensor(named: name) + guard tensor.shape.first.map({ row >= 0 && row < $0 }) == true else { + throw M1Error.invalid("Axis-zero index out of range for \(name)") + } + let slice = try tensor.axisZeroSlice(row) + let pointer = try mappedFile.pointer(to: slice.byteRange) + return try copyTensor(descriptor: slice, pointer: pointer) + } +} diff --git a/ios/Sources/Edge0MLX/OfficialGatedDelta.swift b/ios/Sources/Edge0MLX/OfficialGatedDelta.swift new file mode 100644 index 0000000..acaa296 --- /dev/null +++ b/ios/Sources/Edge0MLX/OfficialGatedDelta.swift @@ -0,0 +1,114 @@ +import MLX + +/// Language-level Swift port of mlx-lm's vector-gated delta Metal kernel used +/// by the official Edge0 8B KDA prefill path. The Metal body, grid and +/// template parameters intentionally retain the upstream structure. +public enum OfficialGatedDelta { + private static let vectorKernel = MLXFast.metalKernel( + name: "edge0_gated_delta_step_vec", + inputNames: ["q", "k", "v", "g", "beta", "state_in", "T"], + outputNames: ["y", "state_out"], + source: """ + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + constexpr int n_per_t = Dk / 32; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + y += b_idx * T * Hv * Dv + hv_idx * Dv; + + auto dk_idx = thread_position_in_threadgroup.x; + auto dv_idx = thread_position_in_grid.y; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + float state[n_per_t]; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = static_cast(i_state[s_idx]); + } + + // g: [B, T, Hv, Dk] + auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk; + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) { + 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_[s_idx]; + kv_mem += state[i] * k_[s_idx]; + } + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + float out = 0.0f; + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + state[i] = state[i] + k_[s_idx] * delta; + out += state[i] * q_[s_idx]; + } + out = simd_sum(out); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv * Dk; + beta_ += Hv; + } + for (int i = 0; i < n_per_t; ++i) { + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + } + """) + + public static func callAsFunction(q: MLXArray, k: MLXArray, v: MLXArray, + decay: MLXArray, beta: MLXArray, + state: MLXArray) throws + -> (output: MLXArray, state: MLXArray) { + guard q.ndim == 4, k.shape == q.shape, v.ndim == 4, + decay.shape == q.shape, + beta.shape == [q.dim(0), q.dim(1), v.dim(2)], + q.dim(2) > 0, v.dim(2) % q.dim(2) == 0, + q.dim(3) > 0, q.dim(3) % 32 == 0, + state.shape == [q.dim(0), v.dim(2), v.dim(3), q.dim(3)], + q.dtype == k.dtype, q.dtype == v.dtype, + q.dtype == decay.dtype, q.dtype == state.dtype else { + throw M1Error.invalid("Invalid official gated-delta tensor shapes") + } + let batch = q.dim(0) + let tokens = q.dim(1) + let keyHeads = q.dim(2) + let keyDimension = q.dim(3) + let valueHeads = v.dim(2) + let valueDimension = v.dim(3) + let outputs = vectorKernel( + [q, k, v, decay, beta, state, tokens], + template: [ + ("InT", q.dtype), ("Dk", keyDimension), + ("Dv", valueDimension), ("Hk", keyHeads), + ("Hv", valueHeads), + ], + grid: (32, valueDimension, batch * valueHeads), + threadGroup: (32, 4, 1), + outputShapes: [ + [batch, tokens, valueHeads, valueDimension], + state.shape, + ], + outputDTypes: [q.dtype, q.dtype]) + return (outputs[0], outputs[1]) + } +} diff --git a/ios/Sources/Edge0MLX/ParallelLoRA.swift b/ios/Sources/Edge0MLX/ParallelLoRA.swift new file mode 100644 index 0000000..ab9c34d --- /dev/null +++ b/ios/Sources/Edge0MLX/ParallelLoRA.swift @@ -0,0 +1,62 @@ +import Edge0Core +import MLX + +/// Parallel, unmerged LoRA used by the official Edge0 deployment. +/// Base INT4 weights stay untouched; fp16 rank-16 adapters are added at run time. +struct ParallelLinear { + let base: QuantizedExpertLinear + private let loraA: MLXArray? + private let loraB: MLXArray? + private let scale: Float + + init(base: QuantizedExpertLinear, loraStore: ExpertTensorStore? = nil, + prefix: String? = nil, rank: Int = 16, alpha: Float = 32) throws { + self.base = base + scale = alpha / Float(rank) + if let loraStore, let prefix { + let a = try loraStore.mlxArrayCopying(named: prefix + ".lora_A") + let b = try loraStore.mlxArrayCopying(named: prefix + ".lora_B") + guard a.shape == [rank, base.inputSize], + b.shape == [base.outputSize, rank] else { + throw M1Error.invalid("LoRA dimensions disagree for \(prefix)") + } + loraA = a + loraB = b + eval(a, b) + } else { + loraA = nil + loraB = nil + } + } + + var inputSize: Int { base.inputSize } + var outputSize: Int { base.outputSize } + + func callAsFunction(_ input: MLXArray) throws -> MLXArray { + let output = try base(input) + guard let loraA, let loraB else { return output } + let x = input.asType(loraA.dtype) + let delta = matmul(matmul(x, loraA.T), loraB.T) * scale + return output + delta.asType(output.dtype) + } +} + +struct ParallelExpert { + let up: ParallelLinear + let gate: ParallelLinear + let down: ParallelLinear + + init(base: QuantizedExpert, loraStore: ExpertTensorStore?, prefix: String) throws { + up = try ParallelLinear(base: base.up, loraStore: loraStore, + prefix: prefix + ".up_proj") + gate = try ParallelLinear(base: base.gate, loraStore: loraStore, + prefix: prefix + ".gate_proj") + down = try ParallelLinear(base: base.down, loraStore: loraStore, + prefix: prefix + ".down_proj") + } + + func callAsFunction(_ input: MLXArray) throws -> MLXArray { + let g = try gate(input) + return try down((g * sigmoid(g)) * up(input)) + } +} diff --git a/ios/Sources/Edge0MLX/QuantizedExpert.swift b/ios/Sources/Edge0MLX/QuantizedExpert.swift new file mode 100644 index 0000000..84c6b60 --- /dev/null +++ b/ios/Sources/Edge0MLX/QuantizedExpert.swift @@ -0,0 +1,227 @@ +import Edge0Core +import Foundation +import MLX + +public enum M1Error: Error, LocalizedError { + case invalid(String) + public var errorDescription: String? { + switch self { case .invalid(let message): return message } + } +} + +/// Affine INT4: W[o,i] = nibble[o,i] * scale[o,i/64] + bias[o,i/64]. +/// Production affine INT4 linear. Coefficients retain their checkpoint dtype +/// and the output follows the activation dtype, matching upstream MLX. +public struct QuantizedExpertLinear { + public let weight: MLXArray + public let scales: MLXArray + public let biases: MLXArray + public let inputSize: Int + public let outputSize: Int + + public init(weight: MLXArray, scales: MLXArray, biases: MLXArray) throws { + guard weight.ndim == 2, weight.dtype == .uint32, + weight.dim(0) > 0, weight.dim(1) > 0, weight.dim(1) % 8 == 0, + scales.shape == [weight.dim(0), weight.dim(1) / 8], + biases.shape == scales.shape, + [.float32, .float16, .bfloat16].contains(scales.dtype), + [.float32, .float16, .bfloat16].contains(biases.dtype) else { + throw M1Error.invalid("Expected affine INT4 [out,in/8] U32 and floating [out,in/64] coefficients") + } + self.weight = weight + self.scales = scales + self.biases = biases + inputSize = weight.dim(1) * 8 + outputSize = weight.dim(0) + } + + public func callAsFunction(_ input: MLXArray) throws -> MLXArray { + guard input.ndim == 2, input.dim(0) > 0, input.dim(1) == inputSize, + [.float32, .float16, .bfloat16].contains(input.dtype) else { + throw M1Error.invalid("Linear input must be floating [tokens,\(inputSize)]") + } + return quantizedMM(input, weight, scales: scales, + biases: biases, transpose: true, groupSize: 64, bits: 4, mode: .affine) + } +} + +public struct QuantizedExpert { + public let up: QuantizedExpertLinear + public let gate: QuantizedExpertLinear + public let down: QuantizedExpertLinear + + public init(up: QuantizedExpertLinear, gate: QuantizedExpertLinear, + down: QuantizedExpertLinear) throws { + guard up.inputSize == gate.inputSize, up.outputSize == gate.outputSize, + down.inputSize == up.outputSize, down.outputSize == up.inputSize else { + throw M1Error.invalid("Inconsistent SwiGLU projection dimensions") + } + self.up = up; self.gate = gate; self.down = down + } + + public func callAsFunction(_ input: MLXArray) throws -> MLXArray { + let g = try gate(input) + return try down((g * sigmoid(g)) * up(input)) + } +} + +/// Official Edge0 whole-layer prefill representation: every routed expert for +/// one projection remains stacked on axis zero and is selected by +/// `gatherQuantizedMM`. Unlike the one-expert correctness path, checkpoint +/// BF16 scales/biases stay BF16, matching `StreamingSwitchGLU.load_full_layer`. +public struct QuantizedExpertStackedLinear { + public let weight: MLXArray + public let scales: MLXArray + public let biases: MLXArray + public let expertCount: Int + public let inputSize: Int + public let outputSize: Int + + public init(weight: MLXArray, scales: MLXArray, biases: MLXArray) throws { + guard weight.ndim == 3, weight.dtype == .uint32, + weight.dim(0) > 0, weight.dim(1) > 0, + weight.dim(2) > 0, weight.dim(2) % 8 == 0, + scales.shape == [weight.dim(0), weight.dim(1), weight.dim(2) / 8], + biases.shape == scales.shape, + [.float32, .float16, .bfloat16].contains(scales.dtype), + [.float32, .float16, .bfloat16].contains(biases.dtype) else { + throw M1Error.invalid("Expected stacked affine INT4 expert tensors") + } + self.weight = weight + self.scales = scales + self.biases = biases + expertCount = weight.dim(0) + inputSize = weight.dim(2) * 8 + outputSize = weight.dim(1) + } + + public func callAsFunction(_ input: MLXArray, expertIndices: MLXArray, + sortedIndices: Bool = true) throws -> MLXArray { + guard input.ndim >= 2, input.dim(-1) == inputSize, + [.uint8, .uint16, .uint32, .uint64, + .int8, .int16, .int32, .int64].contains(expertIndices.dtype), + [.float32, .float16, .bfloat16].contains(input.dtype) else { + throw M1Error.invalid("Invalid stacked expert input or indices") + } + return gatherQuantizedMM( + input, weight, scales: scales, biases: biases, + rhsIndices: expertIndices, transpose: true, + groupSize: 64, bits: 4, mode: .affine, + sortedIndices: sortedIndices) + } +} + +public struct QuantizedExpertStack { + public let up: QuantizedExpertStackedLinear + public let gate: QuantizedExpertStackedLinear + public let down: QuantizedExpertStackedLinear + + public init(up: QuantizedExpertStackedLinear, gate: QuantizedExpertStackedLinear, + down: QuantizedExpertStackedLinear) throws { + guard up.expertCount == gate.expertCount, + up.expertCount == down.expertCount, + up.inputSize == gate.inputSize, + up.outputSize == gate.outputSize, + down.inputSize == up.outputSize, + down.outputSize == up.inputSize else { + throw M1Error.invalid("Inconsistent stacked SwiGLU dimensions") + } + self.up = up + self.gate = gate + self.down = down + } + + /// Official staged-slot assembly: selected expert bundles are stacked in + /// prediction order, then consumed by gathered quantized matmuls. + public init(experts: [QuantizedExpert]) throws { + guard !experts.isEmpty else { + throw M1Error.invalid("Cannot stack an empty expert selection") + } + func projection(_ keyPath: KeyPath) throws + -> QuantizedExpertStackedLinear { + let values = experts.map { $0[keyPath: keyPath] } + return try QuantizedExpertStackedLinear( + weight: stacked(values.map(\.weight), axis: 0), + scales: stacked(values.map(\.scales), axis: 0), + biases: stacked(values.map(\.biases), axis: 0)) + } + try self.init(up: projection(\.up), gate: projection(\.gate), + down: projection(\.down)) + } +} + +public extension ExpertTensorStore { + /// Copies one named tensor, used for the small router and shared expert. + func mlxArrayCopying(named name: String) throws -> MLXArray { + let descriptor = try index.tensor(named: name) + let pointer = try mappedFile.pointer(to: descriptor.byteRange) + return try copyTensor(descriptor: descriptor, pointer: pointer) + } + + func loadExpert(layer: Int, expert: Int) throws -> QuantizedExpert { + guard (0.. QuantizedExpertLinear { + try QuantizedExpertLinear( + weight: mlxArrayCopying(layer: layer, expert: expert, projection: p, part: .weight), + scales: mlxArrayCopying(layer: layer, expert: expert, projection: p, part: .scales), + biases: mlxArrayCopying(layer: layer, expert: expert, projection: p, part: .biases)) + } + return try QuantizedExpert(up: projection(.up), gate: projection(.gate), down: projection(.down)) + } + + func loadSharedExpert(layer: Int) throws -> QuantizedExpert { + func projection(_ p: ExpertProjection) throws -> QuantizedExpertLinear { + let prefix = "model.layers.\(layer).mlp.shared_experts.\(p.rawValue)" + return try QuantizedExpertLinear(weight: mlxArrayCopying(named: prefix + ".weight"), + scales: mlxArrayCopying(named: prefix + ".scales"), + biases: mlxArrayCopying(named: prefix + ".biases")) + } + return try QuantizedExpert(up: projection(.up), gate: projection(.gate), down: projection(.down)) + } + + func loadDenseMLP(layer: Int) throws -> QuantizedExpert { + func projection(_ p: ExpertProjection) throws -> QuantizedExpertLinear { + let prefix = "model.layers.\(layer).mlp.\(p.rawValue)" + return try QuantizedExpertLinear(weight: mlxArrayCopying(named: prefix + ".weight"), + scales: mlxArrayCopying(named: prefix + ".scales"), + biases: mlxArrayCopying(named: prefix + ".biases")) + } + return try QuantizedExpert(up: projection(.up), gate: projection(.gate), down: projection(.down)) + } + + /// Direct translation of official `StreamingSwitchGLU.load_full_layer` for + /// Edge0-8B's separate gate/up/down layout: nine whole stacked tensors. + func loadFullExpertLayer(layer: Int) throws -> QuantizedExpertStack { + func projection(_ p: ExpertProjection) throws -> QuantizedExpertStackedLinear { + let prefix = "model.layers.\(layer).mlp.experts.\(p.rawValue)." + return try QuantizedExpertStackedLinear( + weight: mlxArrayCopying(named: prefix + "weight"), + scales: mlxArrayCopying(named: prefix + "scales"), + biases: mlxArrayCopying(named: prefix + "biases")) + } + return try QuantizedExpertStack(up: projection(.up), gate: projection(.gate), + down: projection(.down)) + } +} + +func copyTensor(descriptor: TensorDescriptor, pointer: UnsafeMutableRawPointer) throws -> MLXArray { + let dtype: DType + switch descriptor.dtype { + case .u32: dtype = .uint32 + case .bf16: dtype = .bfloat16 + case .f16: dtype = .float16 + case .f32: dtype = .float32 + default: throw Edge0MLXError.unsupportedDType(descriptor.dtype) + } + var elements = 1 + for dimension in descriptor.shape { + let product = elements.multipliedReportingOverflow(by: dimension) + guard dimension > 0, !product.overflow else { throw M1Error.invalid("Invalid tensor shape") } + elements = product.partialValue + } + let size = elements.multipliedReportingOverflow(by: descriptor.dtype.byteWidth) + guard !size.overflow, size.partialValue == descriptor.byteCount else { + throw M1Error.invalid("Tensor shape/byte count mismatch: \(descriptor.name)") + } + return MLXArray(Data(bytes: pointer, count: descriptor.byteCount), descriptor.shape, dtype: dtype) +} diff --git a/ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift b/ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift new file mode 100644 index 0000000..06a6bfd --- /dev/null +++ b/ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift @@ -0,0 +1,295 @@ +import Edge0Core +import Foundation +import MLX + +public struct StreamingModelOutput { + public let logits: MLXArray + public let activeMemory: Int + public let peakMemory: Int + public let cacheMemory: Int +} + +private enum StreamingLayer { + case kda(StreamingKDADecoderBlock, KDAStateCache) + case mla(StreamingMLADecoderBlock, MLAKVCache) + + func callAsFunction(_ hidden: MLXArray, + predictedRoute: PredictedMoERoute? = nil) throws -> MLXArray { + switch self { + case .kda(let block, let cache): + return try block(hidden, cache: cache, predictedRoute: predictedRoute).hidden + case .mla(let block, let cache): + return try block(hidden, position: cache.count, cache: cache, + predictedRoute: predictedRoute).hidden + } + } + + var prerouterFeature: LingPrerouterFeature? { + switch self { + case .kda(let block, _): + guard let input = block.prerouterInput, + let indices = block.prerouterIndices else { return nil } + return LingPrerouterFeature(input: input, routeIndices: indices) + case .mla(let block, _): + guard let input = block.prerouterInput, + let indices = block.prerouterIndices else { return nil } + return LingPrerouterFeature(input: input, routeIndices: indices) + } + } + + func stage(experts: [Int]) throws { + switch self { + case .kda(let block, _): try block.stage(experts: experts) + case .mla(let block, _): try block.stage(experts: experts) + } + } + + func reset() { + switch self { + case .kda(let block, let cache): + cache.reset(); block.resetStaging() + case .mla(let block, let cache): + cache.reset(); block.resetStaging() + } + } + + func prefill(_ hidden: MLXArray) throws -> MLXArray { + switch self { + case .kda(let block, let cache): + return try block.prefill(hidden, cache: cache).hidden + case .mla(let block, let cache): + return try block.prefill( + hidden, position: cache.count, cache: cache).hidden + } + } + + func loadFullExpertLayer() throws { + switch self { + case .kda(let block, _): try block.loadFullExpertLayer() + case .mla(let block, _): try block.loadFullExpertLayer() + } + } + + func clearFullExpertLayer() { + switch self { + case .kda(let block, _): block.clearFullExpertLayer() + case .mla(let block, _): block.clearFullExpertLayer() + } + } + + var moeStats: StreamingMoEStats? { + switch self { + case .kda(let block, _): block.moeStats + case .mla(let block, _): block.moeStats + } + } +} + +public struct StreamingRuntimeStats: Sendable { + public let layers: [StreamingMoEStats] + public let cacheHits: Int + public let cacheMisses: Int + public let cachedExperts: Int +} + +/// One-token exact-streaming Edge0 model. +/// +/// Resident tensors are limited to the non-routed backbone, routers, and shared +/// experts. Routed experts remain in the checkpoint mmap and are copied only for +/// the selected top-8, matching the official repository's memory contract. MLX's +/// reusable free-buffer cache is capped at 256 MiB, also matching the official +/// engine default. +public final class StreamingEdge0Model8B { + private let configuration: Edge0Configuration8B + private let store: ExpertTensorStore + private let layers: [StreamingLayer] + private let finalNorm: MLXArray + private let lmHead: QuantizedExpertLinear + private let prerouter: LingPrerouter? + private let bundleCache: ExpertBundleCache + private var predictedRoutes: [Int: PredictedMoERoute] = [:] + + public init(configuration c: Edge0Configuration8B, store: ExpertTensorStore, + cacheLimitBytes: Int = 256 * 1024 * 1024, + loraWeightsURL: URL? = nil, + prerouterWeightsURL: URL? = nil, + progress: (String) -> Void = { _ in }) throws { + guard c.quantization.bits == 4, c.quantization.groupSize == 64, + c.quantization.mode == "affine" else { + throw M1Error.invalid("Full model requires affine INT4/group-64 weights") + } + Memory.cacheLimit = cacheLimitBytes + configuration = c + self.store = store + finalNorm = try store.mlxArrayCopying(named: "model.norm.weight") + lmHead = try QuantizedExpertLinear( + weight: store.mlxArrayCopying(named: "lm_head.weight"), + scales: store.mlxArrayCopying(named: "lm_head.scales"), + biases: store.mlxArrayCopying(named: "lm_head.biases")) + guard finalNorm.shape == [c.hiddenSize], lmHead.inputSize == c.hiddenSize, + lmHead.outputSize == c.vocabularySize else { + throw M1Error.invalid("Final norm or lm-head dimensions disagree with config") + } + progress("final norm and lm-head loaded") + + let bundleCache = ExpertBundleCache(capacity: 64) + self.bundleCache = bundleCache + let loraStore = try loraWeightsURL.map { + try ExpertTensorStore(modelURL: $0, expertCount: c.numExperts) + } + var built: [StreamingLayer] = [] + built.reserveCapacity(c.hiddenLayers) + for layer in 0.. Void = { _, _ in }) throws -> StreamingModelOutput { + let c = configuration + guard tokenID >= 0, tokenID < c.vocabularySize else { + throw M1Error.invalid("Token ID out of vocabulary range") + } + var hidden = try embedding(tokenID: tokenID) + var features: [Int: LingPrerouterFeature] = [:] + for (index, layer) in layers.enumerated() { + hidden = try layer(hidden, predictedRoute: predictedRoutes[index]) + if (7...22).contains(index), let feature = layer.prerouterFeature { + features[index] = feature + } + afterLayer(index, hidden) + } + hidden = MLXFast.rmsNorm(hidden, weight: finalNorm, eps: c.rmsNormEps) + let logits = try lmHead(hidden) + eval(logits) + try updateStagedRoutes(features) + let memory = Memory.snapshot() + return StreamingModelOutput(logits: logits, + activeMemory: memory.activeMemory, + peakMemory: memory.peakMemory, + cacheMemory: memory.cacheMemory) + } + + /// Official chunked prefill path: each sparse layer's complete expert + /// stack is loaded immediately before that layer and dropped when the + /// next layer begins. KDA and MLA process the full token chunk at once. + public func prefill(tokenIDs: [Int], + afterLayer: (Int, MLXArray) -> Void = { _, _ in }) throws + -> StreamingModelOutput { + let c = configuration + guard tokenIDs.count > 1, + tokenIDs.allSatisfy({ (0..