Spaces:
Sleeping
Sleeping
Tomas Marques commited on
Commit ·
2b41258
0
Parent(s):
Latest changes (squashed)
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .claude/settings.json +17 -0
- .gitignore +37 -0
- CLAUDE.md +445 -0
- FlashMoE-iOS/.gitignore +6 -0
- FlashMoE-iOS/397B_ANALYSIS.md +156 -0
- FlashMoE-iOS/App/AppDelegate.swift +22 -0
- FlashMoE-iOS/App/FlashMoEApp.swift +27 -0
- FlashMoE-iOS/Bridge/FlashMoE-Bridging-Header.h +5 -0
- FlashMoE-iOS/Bridge/FlashMoEBridge.swift +419 -0
- FlashMoE-iOS/FlashMoE.xcodeproj/project.pbxproj +516 -0
- FlashMoE-iOS/FlashMoEEngine/FlashMoEEngine.h +157 -0
- FlashMoE-iOS/FlashMoEEngine/FlashMoEEngine.m +1372 -0
- FlashMoE-iOS/IOS_PORT.md +347 -0
- FlashMoE-iOS/Models/ModelCatalog.swift +141 -0
- FlashMoE-iOS/REVIEW.md +206 -0
- FlashMoE-iOS/Resources/Assets.xcassets/AppIcon.appiconset/Contents.json +13 -0
- FlashMoE-iOS/Resources/Assets.xcassets/Contents.json +6 -0
- FlashMoE-iOS/Resources/FlashMoE.entitlements +16 -0
- FlashMoE-iOS/Services/DownloadManager.swift +512 -0
- FlashMoE-iOS/Views/BenchmarkView.swift +547 -0
- FlashMoE-iOS/Views/ChatView.swift +538 -0
- FlashMoE-iOS/Views/ContentView.swift +22 -0
- FlashMoE-iOS/Views/ModelDownloadRow.swift +232 -0
- FlashMoE-iOS/Views/ModelListView.swift +1212 -0
- FlashMoE-iOS/Views/ProfilerView.swift +274 -0
- FlashMoE-iOS/project.yml +75 -0
- README.md +1 -0
- app.py +164 -0
- autoresearch/baseline.txt +1 -0
- autoresearch/benchmark.sh +121 -0
- autoresearch/experiments.tsv +11 -0
- autoresearch/findings.md +35 -0
- autoresearch/prepare.sh +103 -0
- autoresearch/program.md +162 -0
- autoresearch/program_v2.md +212 -0
- autoresearch/research_queue.md +34 -0
- autoresearch/seed_papers.md +79 -0
- build_expert_index.py +167 -0
- cloud_repack.py +1355 -0
- copy_model_to_iphone.sh +140 -0
- docs/context-optimization.md +137 -0
- docs/expert-settings-guide.md +125 -0
- docs/io-and-gpu-exploration.md +215 -0
- docs/ios-port.md +244 -0
- docs/oom-prevention.md +136 -0
- docs/optimization-experiments-q4.md +155 -0
- docs/plan-async-pread-pipeline.md +166 -0
- docs/plan-io-experiments.md +198 -0
- docs/quantization-guide.md +187 -0
- docs/superpowers/plans/2026-03-20-runtime-model-config.md +847 -0
.claude/settings.json
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"permissions": {
|
| 3 |
+
"allow": [
|
| 4 |
+
"Bash(xargs ls:*)",
|
| 5 |
+
"Bash(xcrun devicectl:*)",
|
| 6 |
+
"Read(//Users/anemll/Models/flash/qwen3.5-35b-a3b-tiered/**)",
|
| 7 |
+
"Bash(DEVICE=\"00008150-0016405C01F0401C\" BUNDLE=\"flashmoe.anemll.com\" SRC=\"/Users/anemll/Models/flash/qwen3.5-35b-a3b-tiered\" DST=\"Documents/qwen3.5-35b-a3b-tiered\" __NEW_LINE_c691a4522ac44477__ echo \"Copying model_weights.bin \\(1.3GB\\)...\" xcrun devicectl device copy to --device \"$DEVICE\" --domain-type appDataContainer --domain-identifier \"$BUNDLE\" --source \"$SRC/model_weights.bin\" --destination \"$DST/model_weights.bin\")",
|
| 8 |
+
"Bash(chmod +x /Users/anemll/SourceRelease/GITHUB/ML_playground/FlashMoE-iOS/copy_model_to_iphone.sh)",
|
| 9 |
+
"Read(//Users/anemll/SourceRelease/GITHUB/ML_playground/FlashMoE-iOS/{for\\(i=1;i<=NF;i++\\) if\\($i ~ /^[0-9A-F]{8}-/**)",
|
| 10 |
+
"Read(//Users/anemll/Library/Containers/com.flashmoe.ios/Data/Documents/qwen3.5-35b-a3b-q4/packed_experts/**)",
|
| 11 |
+
"Bash(gh repo:*)"
|
| 12 |
+
],
|
| 13 |
+
"additionalDirectories": [
|
| 14 |
+
"/tmp"
|
| 15 |
+
]
|
| 16 |
+
}
|
| 17 |
+
}
|
.gitignore
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Build artifacts
|
| 2 |
+
metal_infer/infer
|
| 3 |
+
metal_infer/chat
|
| 4 |
+
metal_infer/metal_infer
|
| 5 |
+
metal_infer/bench_pread
|
| 6 |
+
*.o
|
| 7 |
+
*.air
|
| 8 |
+
*.metallib
|
| 9 |
+
|
| 10 |
+
# Generated model files (recreate with scripts — see CLAUDE.md Quick Start)
|
| 11 |
+
metal_infer/model_weights.bin
|
| 12 |
+
metal_infer/model_weights.json
|
| 13 |
+
metal_infer/vocab.bin
|
| 14 |
+
metal_infer/tokenizer.bin
|
| 15 |
+
|
| 16 |
+
# Python
|
| 17 |
+
__pycache__/
|
| 18 |
+
*.pyc
|
| 19 |
+
.venv/
|
| 20 |
+
*.egg-info/
|
| 21 |
+
*.so
|
| 22 |
+
.python-version
|
| 23 |
+
|
| 24 |
+
# macOS
|
| 25 |
+
.DS_Store
|
| 26 |
+
|
| 27 |
+
# Profiling / experiment artifacts
|
| 28 |
+
*.gputrace/
|
| 29 |
+
*.trace/
|
| 30 |
+
*.npz
|
| 31 |
+
logs/
|
| 32 |
+
*.log
|
| 33 |
+
|
| 34 |
+
# Test binaries and experiment data
|
| 35 |
+
metal_infer/test_lzfse
|
| 36 |
+
metal_infer/repack_experts_lz4
|
| 37 |
+
metal_infer/routing_data.bin
|
CLAUDE.md
ADDED
|
@@ -0,0 +1,445 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Flash-MoE: Running Massive MoE Models on a Laptop
|
| 2 |
+
|
| 3 |
+
> **[Read the paper](paper/flash_moe.pdf)** — Full technical details, 90+ experiments, and the story of how an AI and a human built this in 24 hours.
|
| 4 |
+
|
| 5 |
+
Pure C/Metal inference engine for **Qwen3.5 Mixture-of-Experts** models on Apple Silicon. Runs models from 35B to 397B parameters on machines with as little as 24GB RAM, streaming expert weights from SSD through a custom Metal compute pipeline.
|
| 6 |
+
|
| 7 |
+
No Python runtime. No frameworks. Just C, Objective-C, and hand-tuned Metal shaders. Model architecture is auto-detected from HuggingFace `config.json` — switch models with a single `--model` flag.
|
| 8 |
+
|
| 9 |
+
## Compatible Models
|
| 10 |
+
|
| 11 |
+
Any **Qwen3.5 MoE** model with MLX quantization (`model_type: qwen3_5_moe`) is supported. Use the model manager to discover and download compatible models:
|
| 12 |
+
|
| 13 |
+
| Model | Params | Active | Quant | Disk | Min RAM |
|
| 14 |
+
|-------|--------|--------|-------|------|---------|
|
| 15 |
+
| Qwen3.5-35B-A3B | 35B | 3B | 4-bit | ~18GB | 24GB |
|
| 16 |
+
| Qwen3.5-35B-A3B | 35B | 3B | 8-bit | ~35GB | 48GB |
|
| 17 |
+
| Qwen3.5-122B-A10B | 122B | 10B | 4-bit | ~65GB | 48GB |
|
| 18 |
+
| Qwen3.5-397B-A17B | 397B | 17B | 4-bit | ~209GB | 48GB |
|
| 19 |
+
| Qwen3.5-397B-A17B | 397B | 17B | 6-bit | ~280GB | 64GB |
|
| 20 |
+
| Qwen3.5-397B-A17B | 397B | 17B | 8-bit | ~397GB | 96GB |
|
| 21 |
+
|
| 22 |
+
The engine auto-detects architecture, dimensions, expert counts, quantization, and layer types from `config.json`. No recompilation needed.
|
| 23 |
+
|
| 24 |
+
## Development Workflow (MANDATORY)
|
| 25 |
+
|
| 26 |
+
**NEVER merge feature branches into `develop` without explicit approval.** The flow is:
|
| 27 |
+
|
| 28 |
+
1. **Create `feature/branch-name`** from `develop`
|
| 29 |
+
2. Implement, commit, push to the feature branch
|
| 30 |
+
3. **Notify** the user that it's ready for testing
|
| 31 |
+
4. User tests on device and requests changes if needed
|
| 32 |
+
5. Iterate on the feature branch until the user confirms it works
|
| 33 |
+
6. **Only merge when the user explicitly says "merge it"**
|
| 34 |
+
|
| 35 |
+
**Why:** Merging untested code (batched prefill, fused expert, CMD merge) caused repeated gibberish output and crashes that wasted hours of debugging and device restarts. QA happens on the feature branch, not on develop.
|
| 36 |
+
|
| 37 |
+
**Also:**
|
| 38 |
+
- Always add a **settings toggle** (default OFF for experimental features) so broken features can be disabled without reverting code
|
| 39 |
+
- Before any `git rebase` or destructive git operation, **ask the user first**
|
| 40 |
+
- The `develop` branch should always produce correct output
|
| 41 |
+
- **Bundle ID is `com.alexintosh.flashmoe`** — NEVER change this in project.pbxproj. Changing it creates a new app container on the device, losing all downloaded models (potentially 300GB+).
|
| 42 |
+
|
| 43 |
+
## Results
|
| 44 |
+
|
| 45 |
+

|
| 46 |
+
|
| 47 |
+
### MacBook Pro M3 Max (48GB)
|
| 48 |
+
|
| 49 |
+
| Configuration | tok/s | Quality | Notes |
|
| 50 |
+
|--------------|-------|---------|-------|
|
| 51 |
+
| 4-bit + autoresearch optimizations | **9.7** | Excellent | +34.7% from half-precision x_shared + SIMD + FMA |
|
| 52 |
+
| 4-bit experts, FMA kernel | **4.36** | Excellent | Before autoresearch. Full tool calling. 209GB on disk. |
|
| 53 |
+
| **Tiered (hot=4bit, cold=2bit)** | **4.36+** | **Excellent** | **33% smaller on disk. Auto-detected.** |
|
| 54 |
+
| 2-bit experts, trust OS | 5.74 | Good* | 120GB on disk. *Breaks JSON/tool calling. |
|
| 55 |
+
|
| 56 |
+
### iPhone 17 (12GB, A19)
|
| 57 |
+
|
| 58 |
+
| Model | K | tok/s | Notes |
|
| 59 |
+
|-------|---|-------|-------|
|
| 60 |
+
| Qwen3.5-35B-A3B (4-bit) | 8 | **5.5** | 19.5GB download. Full quality. Full GPU path. |
|
| 61 |
+
| Qwen3.5-35B-A3B (tiered) | 8 | **5.5+** | 13.4GB download. Same quality. |
|
| 62 |
+
| Qwen3.5-397B-A17B (4-bit) | 4 | ~0.003* | *CPU fallback only — Metal 4GB per-buffer limit blocks GPU path. |
|
| 63 |
+
| Qwen3.5-397B-A17B (4-bit) | 4 | ~1-2** | **Projected with split weight files enabling GPU path. |
|
| 64 |
+
|
| 65 |
+
*2-bit quantization produces `\name\` instead of `"name"` in JSON output, making tool calling unreliable. 4-bit is the production configuration.
|
| 66 |
+
|
| 67 |
+
**Tiered mode** keeps frequently-activated experts (top ~25%) at 4-bit quality while requantizing cold experts to 2-bit — reducing disk footprint by ~34% without quality loss. Hot experts are profiled from real workloads. See [docs/tiered-expert-quantization.md](docs/tiered-expert-quantization.md) for the full experiment writeup.
|
| 68 |
+
|
| 69 |
+
## Hardware
|
| 70 |
+
|
| 71 |
+
### Development (MacBook Pro)
|
| 72 |
+
- **Machine**: MacBook Pro, Apple M3 Max
|
| 73 |
+
- **Chip**: 16-core CPU (12P + 4E), 40-core GPU, 16-core ANE
|
| 74 |
+
- **Memory**: 48 GB unified (~400 GB/s bandwidth)
|
| 75 |
+
- **SSD**: 1TB Apple Fabric, **17.5 GB/s sequential read** (measured)
|
| 76 |
+
- **macOS**: 26.2 (Darwin 25.2.0)
|
| 77 |
+
|
| 78 |
+
### Mobile (iPhone 17)
|
| 79 |
+
- **Chip**: A19, ~10-core GPU
|
| 80 |
+
- **Memory**: 12 GB unified
|
| 81 |
+
- **SSD**: ~2.5-3 GB/s NVMe
|
| 82 |
+
- **iOS**: 18+
|
| 83 |
+
|
| 84 |
+
## Architecture
|
| 85 |
+
|
| 86 |
+
Qwen3.5 MoE models use a hybrid attention architecture with GatedDeltaNet (linear attention) and standard full attention layers, each containing a Mixture-of-Experts MLP. Model dimensions, expert counts, and layer types vary per model and are read from `config.json` at startup. For example, the 397B model has 60 layers (45 linear + 15 full), 512 experts (K=4 active), hidden dim 4096; the 35B model has 40 layers (30 linear + 10 full), 256 experts (K=8 active), hidden dim 2048.
|
| 87 |
+
|
| 88 |
+
### Key Techniques
|
| 89 |
+
|
| 90 |
+
1. **SSD Expert Streaming** — Expert weights (209GB at 4-bit) are read from NVMe SSD on demand via parallel `pread()` with GCD dispatch groups. Only the K=4 active experts per layer are loaded (~6.75MB each). The OS page cache manages caching — no custom cache needed ("Trust the OS" principle). Inspired by Apple's "LLM in a Flash" paper.
|
| 91 |
+
|
| 92 |
+
1. **Tiered Expert Quantization** — Expert usage follows a Zipfian distribution: ~25% of experts handle ~80% of activations. Hot experts stay at 4-bit; cold experts are requantized to 2-bit (44% smaller each). This shrinks total expert disk by ~34%, improving OS page cache hit rates without quality degradation. Per-expert Metal kernel dispatch selects the right dequant shader at runtime.
|
| 93 |
+
|
| 94 |
+
2. **FMA-Optimized Dequant Kernel** — The inner loop of the 4-bit dequantized matrix-vector multiply rearranges the math from `(nibble * scale + bias) * x` to `fma(nibble, scale*x, bias*x)`. Pre-computing `scale*x` and `bias*x` lets the GPU fused multiply-add unit do dequant+multiply in one instruction. 12% faster than the naive formulation.
|
| 95 |
+
|
| 96 |
+
3. **Metal Compute Shaders** — Hand-written Metal kernels for:
|
| 97 |
+
- 4-bit and 2-bit dequantized matrix-vector multiply (tiled, SIMD-reduced, shared input cache, FMA-optimized)
|
| 98 |
+
- Fused SwiGLU activation
|
| 99 |
+
- RMS normalization (two-pass: sum-of-squares reduction + apply)
|
| 100 |
+
- Batched GPU attention (Q@K^T, softmax, scores@V) for full attention layers
|
| 101 |
+
- GPU RoPE (fused with Q deinterleave and K normalization)
|
| 102 |
+
- MoE combine + residual + sigmoid gate (fused kernel)
|
| 103 |
+
|
| 104 |
+
4. **Deferred GPU Expert Compute** — CMD3 (expert forward pass) is submitted without waiting. The GPU executes it while the CPU prepares the next layer. The combine + residual + norm are also on GPU, feeding directly into the next layer's attention projections.
|
| 105 |
+
|
| 106 |
+
5. **Accelerate BLAS for Linear Attention** — The GatedDeltaNet recurrence uses `cblas_sscal`, `cblas_sgemv`, and `cblas_sger` for the 64-head × 128×128 state matrix update. 64% faster than scalar code.
|
| 107 |
+
|
| 108 |
+
6. **Trust the OS** — No custom expert cache. The OS page cache (~35GB) manages expert data caching via standard LRU. Every custom caching approach we tested (Metal LRU, malloc cache, LZ4 compressed cache) was slower due to GPU memory pressure or overhead. The page cache achieves ~71% hit rate naturally.
|
| 109 |
+
|
| 110 |
+
7. **K-Reduction for Mobile** — MoE models have a natural inference knob: activate fewer experts per token than trained. K=4 instead of K=10 on the 397B model cuts I/O by 60% per token with graceful quality degradation. The router still picks the *best* K experts from all 512 — it's not random subsampling. This makes the 397B viable on iPhone (12GB RAM, 2.5 GB/s NVMe) at ~1-2 tok/s.
|
| 111 |
+
|
| 112 |
+
8. **Half-Precision Shared Memory** — Storing the threadgroup input cache as `half` instead of `float` in Metal dequant kernels halves shared memory usage (16KB → 8KB), doubling GPU occupancy. Since inputs are already approximate from 4-bit dequantization, the precision loss is negligible. +12% tok/s on the main kernel.
|
| 113 |
+
|
| 114 |
+
9. **iOS Unity Build** — The entire 7,500-line inference engine compiles into the iOS app via `#include "infer.m"`. No fork, no separate codebase. A thin C API (`FlashMoEEngine.h`) wraps the static globals, and a Swift `@Observable` bridge provides `AsyncStream<Token>` generation with automatic memory-adaptive context sizing.
|
| 115 |
+
|
| 116 |
+
10. **FP8 E4M3 KV Cache** — Opt-in quantization of the KV cache from float32 to FP8 E4M3 (1 sign, 4 exponent, 3 mantissa bits). Per-position dynamic scales stored separately. Reduces KV memory from ~60KB/position to ~15KB/position for the 397B model (4x reduction), enabling longer context on memory-constrained devices. GPU inline dequant in the fused attention kernel reads FP8 bytes and scales on the fly. Enabled with `--fp8-kv` flag; default off to preserve float32 precision. Note: FP8/sliding window flags must be set BEFORE `metal_setup()` to ensure correct buffer allocation.
|
| 117 |
+
|
| 118 |
+
11. **Fused Online Softmax Attention** — Single-kernel FlashAttention-style implementation replaces the previous 3-dispatch pipeline (Q@K^T, softmax, scores@V) for full-attention layers. Iterates over KV positions in blocks of `BLOCK_SIZE=64`, maintaining online softmax state (running max `m`, running sum `l`, output accumulator `o`) per head. Each block computes partial scores, updates the running statistics, and rescales the accumulator — never materializing the full attention matrix. Reduces 3 GPU dispatches to 1 per full-attention layer.
|
| 119 |
+
|
| 120 |
+
12. **Metal Function Constants** — Compile-time specialization of the fused attention kernel via Metal `[[function_constant(0)]]`. The `USE_FP8_KV` boolean constant eliminates dead branches at pipeline creation time, so the FP8 dequant path has zero overhead when disabled and the float32 path has zero overhead when FP8 is active. Both variants share a single source kernel (`fused_attention_fc`).
|
| 121 |
+
|
| 122 |
+
13. **OOM Prevention** — Comprehensive allocation hardening across the engine: 30+ static scratch buffers pre-allocated at model load (eliminates ~300 malloc/free per token), all 40+ Metal buffer allocations checked for nil with actionable error messages, `calloc` guards on all CPU allocations with early-return on failure, `posix_memalign` for 2MB-aligned expert I/O buffers with error checking. On iOS: dispatch-source memory pressure handler cancels generation on critical pressure, `didReceiveMemoryWarning` observer as a second line of defense, pre-flight 500MB availability check before every generation call, and adaptive context length sizing via `os_proc_available_memory()`. See [docs/oom-prevention.md](docs/oom-prevention.md).
|
| 123 |
+
|
| 124 |
+
14. **Wired Memory Limit** — Metal's `recommendedMaxWorkingSetSize` API queried at startup and stored in `MetalCtx.recommended_working_set`. Used to constrain KV cache allocation so GPU buffer totals stay within the device's wired memory budget, preventing Metal from evicting buffers to system memory (which causes severe latency spikes).
|
| 125 |
+
|
| 126 |
+
15. **Universal App** — The SwiftUI shell compiles for both iPhone and Mac (via "Designed for iPad" / Mac Catalyst compatibility). Views use `#if os(iOS)` conditional compilation for platform-specific UI (toolbar placement, keyboard dismiss, document picker). The same C inference engine, Metal shaders, and Swift bridge run on both platforms without modification.
|
| 127 |
+
|
| 128 |
+
16. **Sliding Window Attention** — Circular KV buffer for full attention layers. Write position cycles via `cache_pos = kv->len % window_size`. The 30 GatedDeltaNet linear attention layers maintain full context through their 128x128 state matrices (O(1) memory), while only the 10 full attention layers are windowed. With window 4096 + FP8: fixed 40MB KV regardless of conversation length. Enabled via `--sliding-window N` flag.
|
| 129 |
+
|
| 130 |
+
17. **FP16 Accumulation** (experimental) — Optional half-precision accumulation in dequant matvec kernels. Apple's GPU has dedicated fp16 ALUs at 2x throughput. The FMA becomes `fma(half(nibble), half(scale*x), half(bias*x))` with final promotion to float32 via `simd_sum`. Risk: fp16 has ~3 decimal digits of precision; sums of 512+ elements may lose accuracy. Default OFF; toggle in Expert Settings.
|
| 131 |
+
|
| 132 |
+
18. **H2O KV Cache Eviction** (in progress) — Heavy Hitter Oracle eviction for the full-attention KV cache. Tracks cumulative post-softmax attention scores per position. When the cache exceeds the budget, it protects sink tokens (first N, typically 4) and recent tokens (25% of budget), then keeps the highest-scoring "heavy hitter" positions. Compacts both CPU and GPU caches in-place so GPU kernels see a shorter contiguous sequence. Replaces sliding window when both are configured (H2O is strictly better). See [docs/context-optimization.md](docs/context-optimization.md).
|
| 133 |
+
|
| 134 |
+
19. **Custom HuggingFace URL Download** — Users can paste any HuggingFace model URL (e.g. `mlx-community/Qwen3.5-35B-A3B-4bit`) in the iOS/Mac app to resolve and download compatible models not in the built-in catalog. The URL is validated, config.json is fetched to verify compatibility, and the model is added to the download list.
|
| 135 |
+
|
| 136 |
+
20. **macOS Sandbox Entitlements** — The universal app includes sandbox entitlements for file access (`com.apple.security.files.user-selected.read-write`), networking (`com.apple.security.network.client`), extended virtual addressing, and increased memory limits.
|
| 137 |
+
|
| 138 |
+
21. **Paper-Guided Autoresearch v2** — Automated experiment loop that reads the research paper, identifies optimization opportunities, implements them, benchmarks with quality gates, and logs results. See `autoresearch/program_v2.md`.
|
| 139 |
+
|
| 140 |
+
### GPTQ/JANG Quantization Pipeline
|
| 141 |
+
|
| 142 |
+
A 4-phase pipeline for producing high-quality 2-bit experts using GPTQ (Data-aware Weight Quantization) with optional JANG (Jang Adaptive N-bit Grading) mixed-precision assignment. GPTQ uses calibration data to build a Hessian proxy (H = X^T @ X) per expert, then applies blocked column-wise error compensation during quantization. The result: same 2-bit format, but output reconstruction error is dramatically lower than RTN (Round To Nearest). This fixes the broken JSON problem at 2-bit.
|
| 143 |
+
|
| 144 |
+
**Phases:**
|
| 145 |
+
- **Phase 0: MSE-Optimal Clipping** — Grid search over 20 clipping ratios per group of 64 values. 15-30% RMSE reduction. In `repack_experts_2bit.py`.
|
| 146 |
+
- **Phase 1: Calibration Collection** — `--collect-activations` flag dumps expert input vectors. `build_hessian.py` accumulates H = X^T @ X per expert online. `calibrate.sh` runner. 16K tokens minimum.
|
| 147 |
+
- **Phase 2: GPTQ Requantization** — `gptq_requantize.py`. Blocked GPTQ (block_size=128). Automatic fallback to MSE-clip for uncalibrated experts. Safety check: only uses GPTQ if it beats RTN RMSE.
|
| 148 |
+
- **Phase 3: Sensitivity Analysis** — `sensitivity_analysis.py` computes freq x quant_error x layer_weight per expert. Assigns 4-bit to most sensitive experts until target GB budget. `repack_experts_tiered.py` updated with `--gptq-dir` and `--hot-experts` flags.
|
| 149 |
+
|
| 150 |
+
See [docs/quantization-guide.md](docs/quantization-guide.md) for the full technical writeup including the GPTQ algorithm, DWQ vs JANG comparison, and pipeline commands.
|
| 151 |
+
|
| 152 |
+
### Expert Settings UI
|
| 153 |
+
|
| 154 |
+
The iOS/Mac app includes a comprehensive Expert Settings panel with info modals for every toggle. Each setting has an analogy (plain-language explanation) and technical details. The UI uses a compact layout with an info icon to the left of each label. Settings include: Active Experts (K), I/O Fanout, CMD1+CMD2 Merge, Fused Attention, Fused Expert Kernel, Expert Prefetch, FP16 Accumulation, FP8 KV Cache, Max Context Length (4K-32K), Sliding Window, Thinking Mode, and H2O Budget (coming soon). Max generation tokens bumped to 2048. See [docs/expert-settings-guide.md](docs/expert-settings-guide.md).
|
| 155 |
+
|
| 156 |
+
### Model Management
|
| 157 |
+
|
| 158 |
+
- Downloaded models are hidden from the download catalog list (no duplicate entries)
|
| 159 |
+
- Trash icon removed from download catalog rows
|
| 160 |
+
- Custom HuggingFace URL download support
|
| 161 |
+
- Import, export, and delete models on-device
|
| 162 |
+
|
| 163 |
+
### iOS-Specific Constraints
|
| 164 |
+
|
| 165 |
+
- **Metal 4GB per-buffer limit** — iOS Metal buffers cannot exceed 4096 MB, regardless of entitlements. The 35B model (2.5GB weights) fits in a single buffer. The 397B model (5.5GB weights) does not. Attempted workarounds: two overlapping buffers (OOM), staging buffer with memcpy per dispatch (data corruption from in-flight command buffer aliasing), CPU fallback (works, 6 min/token). Solution: split `model_weights.bin` into two <4GB files at packing time.
|
| 166 |
+
- **K-reduction quality varies by model** — K=2 and K=4 on the 397B produce gibberish/degenerate output. The model was trained with K=10 and needs K=6+ for coherence (untested — needs GPU path). The 35B at default K=8 works perfectly.
|
| 167 |
+
- **Debug build overhead** — Metal API Validation adds ~2GB of `MTLDebugComputeCommandEncoder` proxies, causing OOM on iPhone. Must build Release for on-device testing.
|
| 168 |
+
- **Bundle ID migration** — Switching from personal to paid developer team requires a new bundle ID (Apple takes 24-48h to release old ones). Moving 300GB of model data between app containers requires `UIDocumentPickerViewController` with `.moveToService`.
|
| 169 |
+
- **File Provider Storage penalty** — Models accessed via Files app integration go through the file coordination layer, adding I/O latency to every `pread`. Always import models to the app's own Documents directory.
|
| 170 |
+
- **`isExcludedFromBackup`** — Must be set on all model files to prevent iOS from purging 200GB+ of data during storage pressure events.
|
| 171 |
+
|
| 172 |
+
See [FlashMoE-iOS/IOS_PORT.md](FlashMoE-iOS/IOS_PORT.md) for the full iOS porting story and [FlashMoE-iOS/397B_ANALYSIS.md](FlashMoE-iOS/397B_ANALYSIS.md) for the 397B memory/performance analysis.
|
| 173 |
+
|
| 174 |
+
### Pipeline Per Layer (4.28ms average at 4-bit)
|
| 175 |
+
|
| 176 |
+
```
|
| 177 |
+
CMD3(prev) → CMD1: attention projections + delta-net [1.22ms GPU]
|
| 178 |
+
→ CPU: flush results [0.01ms CPU]
|
| 179 |
+
→ CMD2: o_proj + norm + routing + shared [0.55ms GPU]
|
| 180 |
+
→ CPU: softmax + topK routing [0.003ms]
|
| 181 |
+
→ I/O: parallel pread K=4 experts [2.41ms SSD]
|
| 182 |
+
→ CMD3: expert forward + combine + norm [0.04ms encode, DEFERRED]
|
| 183 |
+
```
|
| 184 |
+
|
| 185 |
+
### Unified Memory Constraint
|
| 186 |
+
|
| 187 |
+
On Apple Silicon, SSD DMA and GPU compute share the same memory controller and cannot be profitably overlapped. The GPU's dequant kernels are bandwidth-saturated at ~418 GiB/s. Even small background SSD DMA causes disproportionate GPU latency spikes through memory controller arbitration. The serial pipeline (GPU → SSD → GPU) is hardware-optimal.
|
| 188 |
+
|
| 189 |
+
## Model Manager
|
| 190 |
+
|
| 191 |
+
The model manager helps you find, download, and validate compatible models:
|
| 192 |
+
|
| 193 |
+
```bash
|
| 194 |
+
# List local models and search HuggingFace for compatible ones
|
| 195 |
+
python model_manager.py
|
| 196 |
+
|
| 197 |
+
# Search HuggingFace only
|
| 198 |
+
python model_manager.py --search
|
| 199 |
+
|
| 200 |
+
# List local models only
|
| 201 |
+
python model_manager.py --local
|
| 202 |
+
|
| 203 |
+
# Download a specific model
|
| 204 |
+
python model_manager.py --download mlx-community/Qwen3.5-35B-A3B-4bit
|
| 205 |
+
|
| 206 |
+
# Check if a local model is compatible
|
| 207 |
+
python model_manager.py --check /path/to/model
|
| 208 |
+
```
|
| 209 |
+
|
| 210 |
+
After downloading, prepare the model for inference:
|
| 211 |
+
|
| 212 |
+
```bash
|
| 213 |
+
MODEL=~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit
|
| 214 |
+
|
| 215 |
+
# 1. Build expert index (maps tensor names to byte offsets)
|
| 216 |
+
python build_expert_index.py --model $MODEL
|
| 217 |
+
|
| 218 |
+
# 2. Pack expert weights into per-layer binary files
|
| 219 |
+
python repack_experts.py --index expert_index.json
|
| 220 |
+
|
| 221 |
+
# 3. Extract non-expert weights into a single mmap-friendly binary
|
| 222 |
+
python metal_infer/extract_weights.py --model $MODEL
|
| 223 |
+
|
| 224 |
+
# 4. Run inference (auto-detects weights in model directory)
|
| 225 |
+
cd metal_infer && ./infer --model $MODEL --prompt "Hello" --tokens 20
|
| 226 |
+
```
|
| 227 |
+
|
| 228 |
+
Pre-packed models are available on HuggingFace (no repacking needed):
|
| 229 |
+
- `alexintosh/Qwen3.5-35B-A3B-Q4-FlashMoE` (19.5 GB, 4-bit)
|
| 230 |
+
- `alexintosh/Qwen3.5-35B-A3B-Q4-Tiered-FlashMoE` (13.4 GB, tiered)
|
| 231 |
+
|
| 232 |
+
### Tiered Expert Quantization (Optional)
|
| 233 |
+
|
| 234 |
+
Reduces expert disk footprint by ~34% by keeping hot experts at 4-bit and requantizing cold experts to 2-bit. Recommended for memory-constrained setups:
|
| 235 |
+
|
| 236 |
+
```bash
|
| 237 |
+
# 1. Profile expert usage (run a few diverse prompts)
|
| 238 |
+
./infer --model <MODEL> --prompt "Explain quantum computing" --tokens 200 --freq 2>&1 | tee /tmp/freq1.txt
|
| 239 |
+
./infer --model <MODEL> --prompt "Write a Python function" --tokens 200 --freq 2>&1 | tee /tmp/freq2.txt
|
| 240 |
+
|
| 241 |
+
# 2. Generate hot expert manifest (80% coverage threshold)
|
| 242 |
+
python profile_experts.py --freq-output /tmp/freq1.txt /tmp/freq2.txt --coverage 0.8
|
| 243 |
+
|
| 244 |
+
# 3. Repack experts (creates packed_experts_tiered/)
|
| 245 |
+
python repack_experts_tiered.py --model <MODEL>
|
| 246 |
+
|
| 247 |
+
# 4. Run with --tiered (or auto-detected if packed_experts_tiered/ exists)
|
| 248 |
+
cd metal_infer && ./infer --model <MODEL> --tiered --prompt "Hello" --tokens 20
|
| 249 |
+
```
|
| 250 |
+
|
| 251 |
+
## Quick Start
|
| 252 |
+
|
| 253 |
+
```bash
|
| 254 |
+
cd metal_infer
|
| 255 |
+
make
|
| 256 |
+
|
| 257 |
+
# Run with a specific model (auto-detects architecture from config.json)
|
| 258 |
+
./infer --model ~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit \
|
| 259 |
+
--prompt "Explain quantum computing" --tokens 100
|
| 260 |
+
|
| 261 |
+
# Or set FLASH_MOE_MODEL to avoid passing --model every time
|
| 262 |
+
export FLASH_MOE_MODEL=~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit
|
| 263 |
+
./infer --prompt "Explain quantum computing" --tokens 100
|
| 264 |
+
|
| 265 |
+
# 2-bit inference (faster but breaks tool calling)
|
| 266 |
+
./infer --prompt "Explain quantum computing" --tokens 100 --2bit
|
| 267 |
+
|
| 268 |
+
# Tiered mode (hot=4-bit, cold=2-bit, auto-detected if packed_experts_tiered/ exists)
|
| 269 |
+
./infer --prompt "Explain quantum computing" --tokens 100 --tiered
|
| 270 |
+
|
| 271 |
+
# Interactive chat with tool calling (start server first, then chat client)
|
| 272 |
+
./infer --serve &
|
| 273 |
+
./chat
|
| 274 |
+
|
| 275 |
+
# Per-layer timing breakdown
|
| 276 |
+
./infer --prompt "Hello" --tokens 20 --timing
|
| 277 |
+
```
|
| 278 |
+
|
| 279 |
+
## Project Structure
|
| 280 |
+
|
| 281 |
+
```
|
| 282 |
+
model_manager.py # Model discovery, download, and compatibility checking
|
| 283 |
+
build_expert_index.py # Build expert_index.json from safetensors (step 1 of packing)
|
| 284 |
+
repack_experts.py # 4-bit expert packing from safetensors (step 2)
|
| 285 |
+
profile_experts.py # Expert frequency profiling → hot_experts.json
|
| 286 |
+
repack_experts_tiered.py # Tiered repacking (hot=4-bit, cold=2-bit)
|
| 287 |
+
progress.py # Results visualization (Q2/Q4 tracks)
|
| 288 |
+
results.tsv # Experiment log (58 experiments)
|
| 289 |
+
|
| 290 |
+
metal_infer/
|
| 291 |
+
infer.m # Unity build entry point (86 lines, #includes all modules)
|
| 292 |
+
config.h # ModelConfig struct, constants, macros (438 lines)
|
| 293 |
+
timing.h # Timing, telemetry, tracking globals (256 lines)
|
| 294 |
+
fp8.h # FP8 E4M3 encode/decode, per-tensor dynamic scale, g_use_fp8_kv flag (119 lines)
|
| 295 |
+
weights.h # Tensor manifest, hash table, mmap, bf16 conversion (205 lines)
|
| 296 |
+
cpu_kernels.h # Vocabulary, tokenizer, CPU compute kernels (387 lines)
|
| 297 |
+
metal_ctx.h # MetalCtx, metal_setup(), buffer management, wired memory query (603 lines)
|
| 298 |
+
gpu_dispatch.h # BatchMatvecSpec, batched GPU matmul, expert forward (721 lines)
|
| 299 |
+
expert_io.h # I/O thread pool, parallel pread, cache (827 lines)
|
| 300 |
+
layer_forward.h # RoPE, KVCache, attention, MoE, fused pipeline, scratch buffers (3068 lines)
|
| 301 |
+
generate.h # Inference loop, sampling, HTTP serve, main() (1717 lines)
|
| 302 |
+
shaders.metal # Metal compute kernels (~1500 lines, includes fused attention + FP8 variants)
|
| 303 |
+
chat.m # Interactive chat TUI with tool calling
|
| 304 |
+
tokenizer.h # C BPE tokenizer (single-header, 449 lines)
|
| 305 |
+
main.m # MoE-only benchmark
|
| 306 |
+
Makefile # Build system
|
| 307 |
+
extract_weights.py # Creates model_weights.bin from safetensors
|
| 308 |
+
repack_experts_2bit.py # 4-bit → 2-bit expert requantization (with MSE-optimal clipping)
|
| 309 |
+
gptq_requantize.py # Blocked GPTQ 2-bit requantization with Hessian-guided error compensation
|
| 310 |
+
build_hessian.py # Online Hessian accumulation (H = X^T @ X) per expert from calibration data
|
| 311 |
+
sensitivity_analysis.py # Expert sensitivity scoring (freq × quant_error × layer_weight) and bit-width assignment
|
| 312 |
+
calibrate.sh # Calibration runner — collects expert activations over diverse prompts
|
| 313 |
+
train_predictor.py # Expert routing prediction analysis
|
| 314 |
+
model_weights.bin # Non-expert weights (model-specific, mmap'd)
|
| 315 |
+
model_weights.json # Tensor manifest
|
| 316 |
+
vocab.bin # Vocabulary for token decoding
|
| 317 |
+
tokenizer.bin # Pre-exported BPE tokenizer data
|
| 318 |
+
|
| 319 |
+
FlashMoE-iOS/ # Native iOS app
|
| 320 |
+
FlashMoEEngine/
|
| 321 |
+
FlashMoEEngine.h # C API (create/load/generate/cancel/reset/destroy)
|
| 322 |
+
FlashMoEEngine.m # Unity build wrapping infer.m for iOS
|
| 323 |
+
Bridge/
|
| 324 |
+
FlashMoEBridge.swift # @Observable async Swift wrapper
|
| 325 |
+
Views/
|
| 326 |
+
ChatView.swift # Streaming chat UI with thinking disclosure
|
| 327 |
+
ModelListView.swift # Model discovery + download catalog
|
| 328 |
+
ModelDownloadRow.swift # Download progress with pause/resume
|
| 329 |
+
ProfilerView.swift # Resource monitoring overlay
|
| 330 |
+
Services/
|
| 331 |
+
DownloadManager.swift # Background URLSession model downloads
|
| 332 |
+
Models/
|
| 333 |
+
ModelCatalog.swift # HuggingFace model registry with K recommendations
|
| 334 |
+
App/
|
| 335 |
+
FlashMoEApp.swift # SwiftUI app entry point
|
| 336 |
+
IOS_PORT.md # Full iOS porting documentation
|
| 337 |
+
397B_ANALYSIS.md # 397B on iPhone: memory, Metal limits, K-reduction quality
|
| 338 |
+
project.yml # XcodeGen config (iOS 18+, iPhone only)
|
| 339 |
+
copy_model_to_iphone.sh # Push models to device over USB (pymobiledevice3)
|
| 340 |
+
|
| 341 |
+
autoresearch/ # Automated experiment loop
|
| 342 |
+
program.md # Agent instructions for autonomous optimization
|
| 343 |
+
program_v2.md # Paper-guided autoresearch v2 instructions
|
| 344 |
+
benchmark.sh # Measurement harness with quality gates
|
| 345 |
+
prepare.sh # Baseline setup
|
| 346 |
+
experiments.tsv # Experiment log
|
| 347 |
+
findings.md # Autoresearch findings and results
|
| 348 |
+
|
| 349 |
+
docs/
|
| 350 |
+
context-optimization.md # FP8 KV + sliding window + H2O context management
|
| 351 |
+
expert-settings-guide.md # All Expert Settings with analogies and technical details
|
| 352 |
+
ios-port.md # iOS port overview
|
| 353 |
+
optimization-experiments-q4.md # Q4 optimization experiments
|
| 354 |
+
vulkan-learnings-plan.md # Vulkan fork analysis (all 4 phases complete)
|
| 355 |
+
oom-prevention.md # OOM prevention architecture
|
| 356 |
+
tiered-expert-quantization.md # Tiered quantization experiment writeup
|
| 357 |
+
quantization-guide.md # DWQ/JANG comparison, GPTQ pipeline, quantization formats
|
| 358 |
+
```
|
| 359 |
+
|
| 360 |
+
## What We Tried (and What Worked)
|
| 361 |
+
|
| 362 |
+
### Autoresearch Wins (Automated Experiment Loop)
|
| 363 |
+
| Approach | Result | Impact |
|
| 364 |
+
|----------|--------|--------|
|
| 365 |
+
| Half-precision x_shared (v3 kernel) | Halve shared mem → 2× occupancy | **+12.1% tok/s** |
|
| 366 |
+
| FMA 2-bit dequant kernel | fma(nibble, scale*x, bias*x) | **+6.2% tok/s** |
|
| 367 |
+
| Half-precision x_shared (2-bit kernel) | Same occupancy trick | **+3.3% tok/s** |
|
| 368 |
+
| SIMD reduction in rms_norm_qk | simd_sum replaces serial loop | **+2.1% tok/s** |
|
| 369 |
+
|
| 370 |
+
### Vulkan Fork Analysis (Phases 1-4, All Complete)
|
| 371 |
+
|
| 372 |
+
Analyzed the [Vulkan fork](https://github.com/fluxism/flash-moe-vulkan) and identified 4 optimization phases. Key finding: GPU linear attention was already implemented in our code. All 4 phases have been completed:
|
| 373 |
+
|
| 374 |
+
| Phase | Optimization | Impact | Status |
|
| 375 |
+
|-------|-------------|--------|--------|
|
| 376 |
+
| 1 | Delta-net kernel fusion (merge pass 2+3 in gated_delta_net_step) | Eliminates ~1M device memory reads/token | **Done** |
|
| 377 |
+
| 2 | CMD1+CMD2 merging for linear attention layers | Saves 2.25-4.5ms/token (45 layers x 1 sync point) | **Done** |
|
| 378 |
+
| 3 | Modular decomposition (8081-line infer.m -> 9 focused modules) | 0% perf, major maintainability | **Done** |
|
| 379 |
+
| 4 | Dynamic SIMD width (`[[threads_per_simdgroup]]` in all dequant kernels) | Future-proofing for non-32 SIMD hardware | **Done** |
|
| 380 |
+
|
| 381 |
+
Full analysis: [docs/vulkan-learnings-plan.md](docs/vulkan-learnings-plan.md)
|
| 382 |
+
|
| 383 |
+
### Kept (Manual)
|
| 384 |
+
| Approach | Result | Impact |
|
| 385 |
+
|----------|--------|--------|
|
| 386 |
+
| FMA dequant kernel | GPU compute -12% | **+12% tok/s** |
|
| 387 |
+
| Trust OS page cache | Deleted Metal LRU → +38% | **Foundational** |
|
| 388 |
+
| GPU combine+norm in CMD3 | Eliminates CPU round-trip | **Pipeline** |
|
| 389 |
+
| BLAS delta-net (Accelerate) | cpu_attn 0.78→0.28ms | **+64% attn** |
|
| 390 |
+
| F_NOCACHE for 2-bit | +3% from avoiding page thrash | **2-bit only** |
|
| 391 |
+
| GPU fused attention (RoPE) | +2% for full-attn layers | **Small** |
|
| 392 |
+
| C BPE tokenizer | 180ms vs 3500ms startup | **20x startup** |
|
| 393 |
+
| Deferred CMD3 execution | GPU/CPU overlap | **Pipeline** |
|
| 394 |
+
| Tiered expert quant (hot=4b, cold=2b) | -34% disk, same quality | **Cache hit rate** |
|
| 395 |
+
|
| 396 |
+
### Discarded (58 experiments, highlights)
|
| 397 |
+
| Approach | Result | Why |
|
| 398 |
+
|----------|--------|-----|
|
| 399 |
+
| LZ4 expert compression | -13% | Decompress overhead > warm cache savings |
|
| 400 |
+
| F_RDADVISE prefetch | net 0% | Unified memory: SSD DMA slows GPU -73% |
|
| 401 |
+
| Temporal expert prediction | -18% | 25% hit rate, SSD bandwidth waste |
|
| 402 |
+
| MLP routing predictor | 31% accuracy | Worse than temporal baseline |
|
| 403 |
+
| GPU LUT dequant kernel | -2% | Indirect register access serializes |
|
| 404 |
+
| GPU private buffer compression | -20% pipeline | Blit cost 4×7MB > matvec savings |
|
| 405 |
+
| Spin-poll GPU wait | -23% | CPU thermal competes with GPU |
|
| 406 |
+
| Expert file clustering | 0% | NVMe ignores scatter at 7MB granularity |
|
| 407 |
+
| dispatch_io | -70% | dispatch_data management overhead |
|
| 408 |
+
| mmap expert files | -5x | Per-page fault overhead on cold data |
|
| 409 |
+
| Speculative early routing | -38% | Cache pollution + overhead |
|
| 410 |
+
| MTP speculative decoding | break-even | MoE I/O scales per-token (unlike dense) |
|
| 411 |
+
|
| 412 |
+
### Discarded (iOS-specific)
|
| 413 |
+
| Approach | Result | Why |
|
| 414 |
+
|----------|--------|-----|
|
| 415 |
+
| Single 5.5GB Metal buffer (397B weights) | Crash | Metal hard limit: 4096 MB per buffer, not configurable |
|
| 416 |
+
| Two overlapping Metal buffers (~3GB each) | OOM kill | Metal tracks ~8GB shared memory on 12GB device |
|
| 417 |
+
| 50MB staging buffer + memcpy per dispatch | Data corruption | In-flight command buffers alias single staging buffer; later memcpys overwrite earlier tensor data before GPU reads |
|
| 418 |
+
| K=2 on 397B (trained K=10) | Gibberish | Only 20% of trained expert capacity fires, output distribution collapses |
|
| 419 |
+
| K=4 on 397B (trained K=10) | Degenerate ("!!!!") | 40% capacity insufficient for 512-expert model |
|
| 420 |
+
| File Provider Storage for model access | +latency | File coordination layer adds overhead to every pread |
|
| 421 |
+
|
| 422 |
+
## Notable Bug Fixes
|
| 423 |
+
|
| 424 |
+
| Bug | Root Cause | Fix |
|
| 425 |
+
|-----|-----------|-----|
|
| 426 |
+
| Metal shader loading on iOS | `newDefaultLibrary` returns nil when shaders.metal in Resources instead of Sources | Fall back to pre-compiled metallib in bundle; move shader to Sources build phase |
|
| 427 |
+
| KV cache OOM (2GB per cache) | `MAX_SEQ_LEN=1M` used for allocation | Adaptive runtime cap via `os_proc_available_memory()` |
|
| 428 |
+
| ARC heap corruption on model switch | MetalCtx `free()` without nil-ing `id<>` fields | Nil all Objective-C fields before `free` |
|
| 429 |
+
| Expert mmap jetsam kills on iOS | mmap'ing 112GB of expert files | Disabled expert mmap on iOS, pread-only |
|
| 430 |
+
| 2-bit auto-detection missing in iOS | iOS load path skipped 2-bit directory check | Added 2-bit auto-detection in `flashmoe_load()` |
|
| 431 |
+
| String format mismatch warnings | `%d` for `size_t`, `%f` for `int` | Corrected format specifiers throughout |
|
| 432 |
+
| MAX_K buffer overflow on 397B | Hardcoded `MAX_K=8`, 397B needs K=10 | Bumped to `MAX_K=16` with runtime cap |
|
| 433 |
+
|
| 434 |
+
## Safety
|
| 435 |
+
|
| 436 |
+
The engine explicitly controls memory:
|
| 437 |
+
- Non-expert weights: model-dependent (e.g., 5.5GB for 397B, ~1.5GB for 35B, mmap'd read-only)
|
| 438 |
+
- Metal scratch buffers: ~200MB (desktop), ~500MB (397B on iPhone with reduced context)
|
| 439 |
+
- Expert data streams from SSD on demand — no full model load required
|
| 440 |
+
- No custom caches. Trust the OS page cache for expert LRU.
|
| 441 |
+
- iOS: adaptive context length via `os_proc_available_memory()`, KV caches sized to fit device
|
| 442 |
+
- Wired memory budget: `recommendedMaxWorkingSetSize` constrains Metal buffer totals
|
| 443 |
+
- OOM prevention: 30+ pre-allocated scratch buffers, 40+ Metal nil checks, calloc guards, posix_memalign checks. iOS adds memory pressure handler, didReceiveMemoryWarning observer, and 500MB pre-flight check. See [docs/oom-prevention.md](docs/oom-prevention.md).
|
| 444 |
+
- FP8 KV cache (opt-in): reduces KV memory 4x for longer context on constrained devices
|
| 445 |
+
- Minimum RAM: 8GB iPhone (35B), 12GB iPhone (397B with K=4), 24GB Mac (35B), 48GB Mac (397B)
|
FlashMoE-iOS/.gitignore
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Build artifacts
|
| 2 |
+
build/
|
| 3 |
+
DerivedData/
|
| 4 |
+
*.xcworkspace
|
| 5 |
+
*.xcuserdata/
|
| 6 |
+
xcuserdata/
|
FlashMoE-iOS/397B_ANALYSIS.md
ADDED
|
@@ -0,0 +1,156 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Running Qwen3.5-397B on iPhone — Technical Analysis
|
| 2 |
+
|
| 3 |
+
## Status Summary
|
| 4 |
+
|
| 5 |
+
| Aspect | Status | Detail |
|
| 6 |
+
|--------|--------|--------|
|
| 7 |
+
| Model loading | Works | Config parsed, weights mmap'd, expert FDs opened |
|
| 8 |
+
| 35B inference (GPU) | Works | 5.5 tok/s, full quality |
|
| 9 |
+
| 397B inference (CPU fallback) | Works but impractical | ~6 min/token, correct output |
|
| 10 |
+
| 397B inference (GPU) | Blocked | 5.5GB weights exceed Metal 4GB per-buffer limit |
|
| 11 |
+
| K-reduction (K=2) | Gibberish | "no manager" + random tokens |
|
| 12 |
+
| K-reduction (K=4) | Degenerate | "!!!!" repeated output |
|
| 13 |
+
| K-reduction (K=6+) | Untested | Needs GPU path to evaluate at reasonable speed |
|
| 14 |
+
|
| 15 |
+
## Critical Issues (Resolved)
|
| 16 |
+
|
| 17 |
+
### 1. MAX_K=8 but 397B needs K=10
|
| 18 |
+
|
| 19 |
+
The engine hardcoded `#define MAX_K 8` in the MetalCtx struct (infer.m). The 397B model's config specifies `num_experts_per_tok=10`. Loading with default K=10 writes past the end of `buf_multi_expert_data[8]` — buffer overflow, instant crash.
|
| 20 |
+
|
| 21 |
+
**Fix applied**: Bumped MAX_K to 16 and added runtime cap `min(K, MAX_K)` as a safety net.
|
| 22 |
+
|
| 23 |
+
### 2. Weight Auto-Detection Bug
|
| 24 |
+
|
| 25 |
+
The K-reduction branch added weight auto-detection in the model directory, but the `default_weights` char array wasn't initialized when `model_path` is NULL. Reading `default_weights[0]` is undefined behavior.
|
| 26 |
+
|
| 27 |
+
**Fix applied**: Initialize `default_weights[0] = '\0'` (and same for manifest/vocab) at declaration.
|
| 28 |
+
|
| 29 |
+
## Critical Issue (Open): Metal 4GB Per-Buffer Limit
|
| 30 |
+
|
| 31 |
+
### The Problem
|
| 32 |
+
|
| 33 |
+
Metal on iOS enforces a **hard 4GB per-buffer limit**. The 397B model's non-expert weights (`model_weights.bin`) are 5.5GB.
|
| 34 |
+
|
| 35 |
+
### What We Tried
|
| 36 |
+
|
| 37 |
+
| Approach | Outcome | Root Cause |
|
| 38 |
+
|----------|---------|-----------|
|
| 39 |
+
| Single Metal buffer (5.5GB) | Metal assertion crash | `buffer must not exceed 4096 MB` — hard platform limit, not configurable |
|
| 40 |
+
| Two overlapping Metal buffers (~3GB each) | OOM kill | Metal internally tracks ~8GB of shared memory. On a 12GB device, iOS kills the app. |
|
| 41 |
+
| 50MB staging buffer (memcpy per dispatch) | **Data corruption** | Single staging buffer shared across N in-flight command buffers. Metal executes all encoders AFTER `commit`, so CPU memcpys for later tensors overwrite data before GPU reads earlier tensors. Fundamental architectural flaw — would need N staging buffers for N concurrent command buffers. |
|
| 42 |
+
| CPU fallback (Accelerate BLAS) | Works, 6 min/token | Weight matmuls run on CPU. Expert forward still uses GPU (small per-expert buffers). Correct output but unusable. |
|
| 43 |
+
|
| 44 |
+
### Why Entitlements Don't Help
|
| 45 |
+
|
| 46 |
+
- `extended-virtual-addressing` — expands total virtual address space, not per-buffer limit
|
| 47 |
+
- `increased-memory-limit` — raises the total memory ceiling for the process, not per-buffer
|
| 48 |
+
|
| 49 |
+
Neither entitlement changes the 4GB per-Metal-buffer hard cap.
|
| 50 |
+
|
| 51 |
+
### The Solution: Split Weight Files
|
| 52 |
+
|
| 53 |
+
Split `model_weights.bin` at the Python packing stage into two files:
|
| 54 |
+
|
| 55 |
+
```
|
| 56 |
+
model_weights_0.bin (~3.0 GB) — layers 0-29 tensors
|
| 57 |
+
model_weights_1.bin (~2.5 GB) — layers 30-59 tensors
|
| 58 |
+
```
|
| 59 |
+
|
| 60 |
+
Each file gets its own Metal buffer (both under 4GB). The tensor manifest (`model_weights.json`) gains a `file_index` field so `infer.m` selects the correct mmap base pointer per tensor.
|
| 61 |
+
|
| 62 |
+
**Engineering work required**:
|
| 63 |
+
1. `extract_weights.py` — emit two files with a configurable split point
|
| 64 |
+
2. `model_weights.json` — add file index per tensor entry
|
| 65 |
+
3. `infer.m` — mmap both files, index tensors to correct buffer
|
| 66 |
+
4. Re-upload 397B model to HuggingFace
|
| 67 |
+
|
| 68 |
+
## Memory Budget Analysis
|
| 69 |
+
|
| 70 |
+
### KV Cache
|
| 71 |
+
|
| 72 |
+
GPU_KV_SEQ was hardcoded to 8192. For the 397B with head_dim=256, num_kv_heads=2:
|
| 73 |
+
|
| 74 |
+
| Component | Formula | Size |
|
| 75 |
+
|-----------|---------|------|
|
| 76 |
+
| KV cache | 15 layers x 2 buffers x seq x 512 x 4B | 504 MB (seq=8192), **126 MB (seq=2048)** |
|
| 77 |
+
| Attn scores | 32 heads x seq x 4B x 15 layers | 15 MB (seq=8192), 4 MB (seq=2048) |
|
| 78 |
+
| Delta-net state | 45 layers x 64 x 128 x 128 x 4B | 188 MB |
|
| 79 |
+
| Delta-net conv | 45 layers x 3 x 8192 x 4B | 4 MB |
|
| 80 |
+
| Expert buffers | MAX_K x 2 x 7MB | 112 MB (MAX_K=8), 224 MB (MAX_K=16) |
|
| 81 |
+
| Working buffers | — | ~50 MB |
|
| 82 |
+
|
| 83 |
+
**Fix applied**: Adaptive context length reduces KV seq to 2048 on iPhone, saving ~378 MB.
|
| 84 |
+
|
| 85 |
+
### Total Memory Budget (iPhone 17, 12GB)
|
| 86 |
+
|
| 87 |
+
```
|
| 88 |
+
Non-expert weights (mmap'd): 5.5 GB (virtual, paged on demand)
|
| 89 |
+
Metal buffers (KV=2048): ~500 MB
|
| 90 |
+
iOS overhead: ~2.0 GB
|
| 91 |
+
--------
|
| 92 |
+
Resident estimate: ~3.0 GB
|
| 93 |
+
Available for expert page cache: ~9.0 GB
|
| 94 |
+
```
|
| 95 |
+
|
| 96 |
+
### GQA Ratio
|
| 97 |
+
|
| 98 |
+
The 397B uses 32 query heads but only 2 KV heads (16:1 GQA ratio). The 35B uses 32:4 (8:1). The attention code handles this via `cfg.num_kv_heads` dynamically with `kv_h = h / (num_attn_heads / num_kv_heads)`. The 16:1 ratio works correctly since it's computed from config values.
|
| 99 |
+
|
| 100 |
+
## Performance Projections
|
| 101 |
+
|
| 102 |
+
### With GPU Path (After Split Weights)
|
| 103 |
+
|
| 104 |
+
| Configuration | Expert I/O | I/O Time | GPU Compute | Total | tok/s |
|
| 105 |
+
|--------------|-----------|----------|-------------|-------|-------|
|
| 106 |
+
| K=10, 4-bit | 10 x 60 x 6.75MB = 4.1 GB | ~1.6s | ~0.5s | ~2.1s | ~0.5 |
|
| 107 |
+
| K=4, 4-bit | 4 x 60 x 6.75MB = 1.6 GB | ~0.65s | ~0.4s | ~1.05s | **~1.0** |
|
| 108 |
+
| K=4, tiered | 4 x 60 x ~4.5MB = 1.1 GB | ~0.43s | ~0.4s | ~0.83s | **~1.2** |
|
| 109 |
+
|
| 110 |
+
iPhone NVMe throughput assumed at ~2.5 GB/s.
|
| 111 |
+
|
| 112 |
+
### Page Cache Effectiveness
|
| 113 |
+
|
| 114 |
+
With ~9 GB free for page cache and 208 GB of expert data:
|
| 115 |
+
|
| 116 |
+
- Cache can hold ~1,330 experts out of 30,720 total (60 layers x 512)
|
| 117 |
+
- Only 240 experts used per token at K=4 (4 x 60 layers)
|
| 118 |
+
- Expert usage follows Zipfian: ~25% of experts handle 80% of activations
|
| 119 |
+
- Expected page cache hit rate: **10-15%** (most reads are cache misses)
|
| 120 |
+
- With tiered (cold=2-bit): smaller expert files, maybe **15-20%** hit rate
|
| 121 |
+
|
| 122 |
+
The 397B on iPhone will be heavily I/O bound. Every optimization to reduce I/O (lower K, tiered, expert compression) matters more than GPU optimization.
|
| 123 |
+
|
| 124 |
+
### Storage Requirements
|
| 125 |
+
|
| 126 |
+
| Configuration | Expert Disk | Total Disk | Min iPhone Storage |
|
| 127 |
+
|--------------|------------|------------|-------------------|
|
| 128 |
+
| K=10, 4-bit | 208 GB | 214 GB | 256 GB |
|
| 129 |
+
| K=4, 4-bit | 208 GB | 214 GB | 256 GB |
|
| 130 |
+
| K=4, tiered | ~140 GB | 146 GB | 256 GB |
|
| 131 |
+
| K=4, 2-bit | ~104 GB | 110 GB | 128 GB |
|
| 132 |
+
|
| 133 |
+
K-reduction does NOT save storage (all 512 experts per layer remain on disk for routing). Only tiered/2-bit quantization reduces disk footprint.
|
| 134 |
+
|
| 135 |
+
## K-Reduction Quality Analysis
|
| 136 |
+
|
| 137 |
+
### Why K=2 and K=4 Fail on 397B
|
| 138 |
+
|
| 139 |
+
The 397B model was trained with K=10 active experts per token. Each expert contributes a weighted fraction of the final hidden state. Reducing K means:
|
| 140 |
+
|
| 141 |
+
- **K=2**: Only 20% of the trained expert capacity fires. The router picks the top 2, but the model's residual stream expects contributions from ~10 experts. The output distribution collapses.
|
| 142 |
+
- **K=4**: 40% of trained capacity. Better, but for a model this large (512 experts, K=10), each expert is specialized enough that missing 60% of them produces degenerate patterns.
|
| 143 |
+
- **K=6+**: Untested. The 35B model (trained K=8) works perfectly at K=8, suggesting models tolerate reduction better when the gap is smaller relative to training K.
|
| 144 |
+
|
| 145 |
+
### Testing Constraints
|
| 146 |
+
|
| 147 |
+
K-reduction quality can only be properly evaluated with the GPU path enabled (split weight files). The CPU fallback at 6 min/token makes quality evaluation impractical — you need hundreds of tokens to judge coherence, which would take hours.
|
| 148 |
+
|
| 149 |
+
## Action Items (Priority Order)
|
| 150 |
+
|
| 151 |
+
1. **Split `model_weights.bin` into two <4GB files** — enables GPU path on iOS
|
| 152 |
+
2. **Test K=6, K=8, K=10 with GPU path** — find minimum viable K for coherent 397B output
|
| 153 |
+
3. **Upload split 397B model to HuggingFace** — make it downloadable in iOS app
|
| 154 |
+
4. **Profile 397B tiered experts** — repack hot=4-bit cold=2-bit, reduce storage to ~140 GB
|
| 155 |
+
5. **Benchmark actual I/O throughput on A19** — validate 2.5 GB/s assumption
|
| 156 |
+
6. **Test thermal behavior** — sustained 397B inference may thermal-throttle the A19
|
FlashMoE-iOS/App/AppDelegate.swift
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#if canImport(UIKit) && !isMacro
|
| 2 |
+
/*
|
| 3 |
+
* AppDelegate.swift — Background URL session event handling
|
| 4 |
+
*
|
| 5 |
+
* Required for URLSession background downloads to reconnect
|
| 6 |
+
* when iOS relaunches the app to deliver completed download events.
|
| 7 |
+
*/
|
| 8 |
+
|
| 9 |
+
import UIKit
|
| 10 |
+
|
| 11 |
+
class AppDelegate: NSObject, UIApplicationDelegate {
|
| 12 |
+
func application(
|
| 13 |
+
_ application: UIApplication,
|
| 14 |
+
handleEventsForBackgroundURLSession identifier: String,
|
| 15 |
+
completionHandler: @escaping () -> Void
|
| 16 |
+
) {
|
| 17 |
+
if identifier == "com.flashmoe.model-download" {
|
| 18 |
+
DownloadManager.shared.backgroundCompletionHandler = completionHandler
|
| 19 |
+
}
|
| 20 |
+
}
|
| 21 |
+
}
|
| 22 |
+
#endif
|
FlashMoE-iOS/App/FlashMoEApp.swift
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* FlashMoEApp.swift — Main app entry point
|
| 3 |
+
*
|
| 4 |
+
* Flash-MoE iOS: Run massive MoE models on iPhone.
|
| 5 |
+
* Targets iOS 18+ (iPhone 15 Pro and later with 8GB+ RAM).
|
| 6 |
+
*/
|
| 7 |
+
|
| 8 |
+
import SwiftUI
|
| 9 |
+
|
| 10 |
+
@main
|
| 11 |
+
struct FlashMoEApp: App {
|
| 12 |
+
#if os(iOS)
|
| 13 |
+
@UIApplicationDelegateAdaptor(AppDelegate.self) var appDelegate
|
| 14 |
+
#endif
|
| 15 |
+
@State private var engine: FlashMoEEngine = .init()
|
| 16 |
+
|
| 17 |
+
var body: some Scene {
|
| 18 |
+
WindowGroup {
|
| 19 |
+
ContentView()
|
| 20 |
+
.environment(engine)
|
| 21 |
+
#if os(macOS)
|
| 22 |
+
.frame(minWidth: 900, minHeight: 600)
|
| 23 |
+
#endif
|
| 24 |
+
}
|
| 25 |
+
}
|
| 26 |
+
}
|
| 27 |
+
|
FlashMoE-iOS/Bridge/FlashMoE-Bridging-Header.h
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* FlashMoE-Bridging-Header.h — Exposes the C engine API to Swift
|
| 3 |
+
*/
|
| 4 |
+
|
| 5 |
+
#import "FlashMoEEngine.h"
|
FlashMoE-iOS/Bridge/FlashMoEBridge.swift
ADDED
|
@@ -0,0 +1,419 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* FlashMoEBridge.swift — Async Swift wrapper for the FlashMoE C engine
|
| 3 |
+
*
|
| 4 |
+
* Provides a Swift-native interface with:
|
| 5 |
+
* - AsyncStream for token-by-token generation
|
| 6 |
+
* - Observable properties for SwiftUI integration
|
| 7 |
+
* - Automatic background thread management
|
| 8 |
+
*/
|
| 9 |
+
|
| 10 |
+
import Foundation
|
| 11 |
+
import Observation
|
| 12 |
+
|
| 13 |
+
// MARK: - Data Types
|
| 14 |
+
|
| 15 |
+
/// Generation result with streaming tokens
|
| 16 |
+
struct GenerationToken {
|
| 17 |
+
let text: String
|
| 18 |
+
let tokenId: Int
|
| 19 |
+
let tokensGenerated: Int
|
| 20 |
+
let tokensPerSecond: Double
|
| 21 |
+
}
|
| 22 |
+
|
| 23 |
+
/// Model information after loading
|
| 24 |
+
struct ModelInfo {
|
| 25 |
+
let name: String
|
| 26 |
+
let numLayers: Int
|
| 27 |
+
let numExperts: Int
|
| 28 |
+
let activeExpertsK: Int
|
| 29 |
+
let hiddenDim: Int
|
| 30 |
+
let vocabSize: Int
|
| 31 |
+
let weightFileBytes: UInt64
|
| 32 |
+
let expertFileBytes: UInt64
|
| 33 |
+
let metalBufferBytes: UInt64
|
| 34 |
+
|
| 35 |
+
var weightFileMB: Double { Double(weightFileBytes) / 1_048_576 }
|
| 36 |
+
var expertFileMB: Double { Double(expertFileBytes) / 1_048_576 }
|
| 37 |
+
var totalSizeMB: Double { weightFileMB + expertFileMB }
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
+
/// Engine state for UI binding
|
| 41 |
+
enum EngineState: Equatable {
|
| 42 |
+
case idle
|
| 43 |
+
case loading
|
| 44 |
+
case ready
|
| 45 |
+
case generating
|
| 46 |
+
case error(String)
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
// MARK: - FlashMoEEngine (Observable)
|
| 50 |
+
|
| 51 |
+
@Observable
|
| 52 |
+
final class FlashMoEEngine: @unchecked Sendable {
|
| 53 |
+
// Observable state for SwiftUI
|
| 54 |
+
private(set) var state: EngineState = .idle
|
| 55 |
+
private(set) var modelInfo: ModelInfo?
|
| 56 |
+
private(set) var tokensPerSecond: Double = 0
|
| 57 |
+
private(set) var tokensGenerated: Int = 0
|
| 58 |
+
private(set) var timeToFirstToken: Double = 0
|
| 59 |
+
|
| 60 |
+
// Private engine state
|
| 61 |
+
private var context: OpaquePointer? // FlashMoEContext*
|
| 62 |
+
private let engineQueue = DispatchQueue(label: "com.flashmoe.engine", qos: .userInitiated)
|
| 63 |
+
private var isGenerating = false
|
| 64 |
+
|
| 65 |
+
init() {}
|
| 66 |
+
|
| 67 |
+
deinit {
|
| 68 |
+
if let ctx = context {
|
| 69 |
+
flashmoe_destroy(ctx)
|
| 70 |
+
}
|
| 71 |
+
}
|
| 72 |
+
|
| 73 |
+
// MARK: - Model Loading
|
| 74 |
+
|
| 75 |
+
/// Load a model from the given path. Runs on a background thread.
|
| 76 |
+
/// Load a model. Set `activeExpertsK` to reduce expert count for large models on small devices.
|
| 77 |
+
/// For example, K=4 on a K=10 model cuts I/O by 60%.
|
| 78 |
+
func loadModel(at path: String, maxContext: Int = 0, thinkBudget: Int = 2048,
|
| 79 |
+
useTiered: Bool = false, activeExpertsK: Int = 0, cacheIOSplit: Int = 1,
|
| 80 |
+
cmdMerge: Bool = true, fusedAttention: Bool = false,
|
| 81 |
+
expertPrefetch: Bool = false, fusedExpert: Bool = true,
|
| 82 |
+
fp16Accumulation: Bool = false, fp8KVCache: Bool = false,
|
| 83 |
+
slidingWindow: Int = 0, h2oBudget: Int = 0,
|
| 84 |
+
ropeScalingMode: Int = 0, ropeScaleFactor: Float = 1.0,
|
| 85 |
+
prefillBatch: Int = 1,
|
| 86 |
+
verbose: Bool = false) async throws {
|
| 87 |
+
guard state != .loading && state != .generating else {
|
| 88 |
+
throw FlashMoEError.busy
|
| 89 |
+
}
|
| 90 |
+
|
| 91 |
+
await MainActor.run { state = .loading }
|
| 92 |
+
|
| 93 |
+
return try await withCheckedThrowingContinuation { continuation in
|
| 94 |
+
engineQueue.async { [weak self] in
|
| 95 |
+
guard let self else {
|
| 96 |
+
continuation.resume(throwing: FlashMoEError.engineDestroyed)
|
| 97 |
+
return
|
| 98 |
+
}
|
| 99 |
+
|
| 100 |
+
// Create context if needed
|
| 101 |
+
if self.context == nil {
|
| 102 |
+
self.context = flashmoe_create()
|
| 103 |
+
}
|
| 104 |
+
guard let ctx = self.context else {
|
| 105 |
+
DispatchQueue.main.async { self.state = .error("Failed to create engine context") }
|
| 106 |
+
continuation.resume(throwing: FlashMoEError.initFailed)
|
| 107 |
+
return
|
| 108 |
+
}
|
| 109 |
+
|
| 110 |
+
// Configure
|
| 111 |
+
var config = FlashMoEConfig()
|
| 112 |
+
let pathCStr = (path as NSString).utf8String
|
| 113 |
+
config.model_path = pathCStr
|
| 114 |
+
config.max_context = Int32(maxContext)
|
| 115 |
+
config.think_budget = Int32(thinkBudget)
|
| 116 |
+
config.use_tiered = useTiered ? 1 : 0
|
| 117 |
+
config.active_experts_k = Int32(activeExpertsK)
|
| 118 |
+
config.cache_io_split = Int32(cacheIOSplit)
|
| 119 |
+
config.cmd_merge = cmdMerge ? 1 : 0
|
| 120 |
+
config.fused_attention = fusedAttention ? 1 : 0
|
| 121 |
+
config.expert_prefetch = expertPrefetch ? 1 : 0
|
| 122 |
+
config.fused_expert = fusedExpert ? 1 : 0
|
| 123 |
+
config.fp16_accumulation = fp16Accumulation ? 1 : 0
|
| 124 |
+
config.fp8_kv_cache = fp8KVCache ? 1 : 0
|
| 125 |
+
config.sliding_window = Int32(slidingWindow)
|
| 126 |
+
config.h2o_budget = Int32(h2oBudget)
|
| 127 |
+
config.rope_scaling_mode = Int32(ropeScalingMode)
|
| 128 |
+
config.rope_scale_factor = ropeScaleFactor
|
| 129 |
+
config.prefill_batch = Int32(prefillBatch)
|
| 130 |
+
config.verbose = verbose ? 1 : 0
|
| 131 |
+
|
| 132 |
+
// Load
|
| 133 |
+
let result = flashmoe_load(ctx, &config)
|
| 134 |
+
if result != 0 {
|
| 135 |
+
let error = String(cString: flashmoe_last_error(ctx))
|
| 136 |
+
DispatchQueue.main.async { self.state = .error(error) }
|
| 137 |
+
continuation.resume(throwing: FlashMoEError.loadFailed(error))
|
| 138 |
+
return
|
| 139 |
+
}
|
| 140 |
+
|
| 141 |
+
// Get stats for model info
|
| 142 |
+
var stats = FlashMoEStats()
|
| 143 |
+
flashmoe_get_stats(ctx, &stats)
|
| 144 |
+
|
| 145 |
+
let modelName = withUnsafePointer(to: &stats.model_name) {
|
| 146 |
+
$0.withMemoryRebound(to: CChar.self, capacity: 256) { String(cString: $0) }
|
| 147 |
+
}
|
| 148 |
+
let info = ModelInfo(
|
| 149 |
+
name: modelName,
|
| 150 |
+
numLayers: Int(stats.num_layers),
|
| 151 |
+
numExperts: Int(stats.num_experts),
|
| 152 |
+
activeExpertsK: Int(stats.active_experts_k),
|
| 153 |
+
hiddenDim: Int(stats.hidden_dim),
|
| 154 |
+
vocabSize: Int(stats.vocab_size),
|
| 155 |
+
weightFileBytes: UInt64(stats.weight_file_bytes),
|
| 156 |
+
expertFileBytes: UInt64(stats.expert_file_bytes),
|
| 157 |
+
metalBufferBytes: UInt64(stats.metal_buffer_bytes)
|
| 158 |
+
)
|
| 159 |
+
|
| 160 |
+
DispatchQueue.main.async {
|
| 161 |
+
self.modelInfo = info
|
| 162 |
+
self.state = .ready
|
| 163 |
+
}
|
| 164 |
+
continuation.resume()
|
| 165 |
+
}
|
| 166 |
+
}
|
| 167 |
+
}
|
| 168 |
+
|
| 169 |
+
/// Unload the current model
|
| 170 |
+
func unloadModel() {
|
| 171 |
+
guard let ctx = context else { return }
|
| 172 |
+
engineQueue.sync {
|
| 173 |
+
flashmoe_unload(ctx)
|
| 174 |
+
}
|
| 175 |
+
modelInfo = nil
|
| 176 |
+
state = .idle
|
| 177 |
+
}
|
| 178 |
+
|
| 179 |
+
// MARK: - Generation
|
| 180 |
+
|
| 181 |
+
/// Generate tokens from a prompt, returning an AsyncStream of tokens
|
| 182 |
+
func generate(prompt: String, maxTokens: Int = 200) -> AsyncStream<GenerationToken> {
|
| 183 |
+
AsyncStream { continuation in
|
| 184 |
+
guard let ctx = context, state == .ready else {
|
| 185 |
+
continuation.finish()
|
| 186 |
+
return
|
| 187 |
+
}
|
| 188 |
+
|
| 189 |
+
DispatchQueue.main.async {
|
| 190 |
+
self.state = .generating
|
| 191 |
+
self.tokensGenerated = 0
|
| 192 |
+
self.tokensPerSecond = 0
|
| 193 |
+
self.isGenerating = true
|
| 194 |
+
}
|
| 195 |
+
|
| 196 |
+
// Set up cancellation
|
| 197 |
+
nonisolated(unsafe) let ctxForCancel = ctx
|
| 198 |
+
continuation.onTermination = { @Sendable _ in
|
| 199 |
+
flashmoe_cancel(ctxForCancel)
|
| 200 |
+
}
|
| 201 |
+
|
| 202 |
+
engineQueue.async { [weak self] in
|
| 203 |
+
// C callback bridge: userdata points to the Swift continuation
|
| 204 |
+
let userDataPtr = Unmanaged.passRetained(
|
| 205 |
+
TokenCallbackContext(continuation: continuation, engine: self)
|
| 206 |
+
).toOpaque()
|
| 207 |
+
|
| 208 |
+
let result = flashmoe_generate(
|
| 209 |
+
ctx,
|
| 210 |
+
prompt,
|
| 211 |
+
Int32(maxTokens),
|
| 212 |
+
{ tokenText, tokenId, tokensGenerated, tokensPerSecond, userData -> Int32 in
|
| 213 |
+
guard let userData else { return 1 }
|
| 214 |
+
let context = Unmanaged<TokenCallbackContext>.fromOpaque(userData)
|
| 215 |
+
.takeUnretainedValue()
|
| 216 |
+
|
| 217 |
+
guard let text = tokenText else { return 0 }
|
| 218 |
+
let token = GenerationToken(
|
| 219 |
+
text: String(cString: text),
|
| 220 |
+
tokenId: Int(tokenId),
|
| 221 |
+
tokensGenerated: Int(tokensGenerated),
|
| 222 |
+
tokensPerSecond: tokensPerSecond
|
| 223 |
+
)
|
| 224 |
+
|
| 225 |
+
// Update engine stats on main thread
|
| 226 |
+
if let engine = context.engine {
|
| 227 |
+
DispatchQueue.main.async {
|
| 228 |
+
engine.tokensGenerated = Int(tokensGenerated)
|
| 229 |
+
engine.tokensPerSecond = tokensPerSecond
|
| 230 |
+
}
|
| 231 |
+
}
|
| 232 |
+
|
| 233 |
+
context.continuation.yield(token)
|
| 234 |
+
return 0
|
| 235 |
+
},
|
| 236 |
+
userDataPtr
|
| 237 |
+
)
|
| 238 |
+
|
| 239 |
+
// Clean up
|
| 240 |
+
Unmanaged<TokenCallbackContext>.fromOpaque(userDataPtr).release()
|
| 241 |
+
|
| 242 |
+
// Get final stats
|
| 243 |
+
var stats = FlashMoEStats()
|
| 244 |
+
flashmoe_get_stats(ctx, &stats)
|
| 245 |
+
|
| 246 |
+
DispatchQueue.main.async {
|
| 247 |
+
self?.timeToFirstToken = stats.ttft_ms
|
| 248 |
+
self?.tokensPerSecond = stats.tokens_per_second
|
| 249 |
+
self?.tokensGenerated = Int(stats.tokens_generated)
|
| 250 |
+
self?.state = .ready
|
| 251 |
+
self?.isGenerating = false
|
| 252 |
+
}
|
| 253 |
+
|
| 254 |
+
continuation.finish()
|
| 255 |
+
}
|
| 256 |
+
}
|
| 257 |
+
}
|
| 258 |
+
|
| 259 |
+
/// Generate continuation — reuses KV cache from previous turns.
|
| 260 |
+
/// Returns nil if context is full (caller should reset and use generate instead).
|
| 261 |
+
func generateContinuation(userMessage: String, maxTokens: Int = 200) -> AsyncStream<GenerationToken> {
|
| 262 |
+
AsyncStream { continuation in
|
| 263 |
+
guard let ctx = context, state == .ready else {
|
| 264 |
+
continuation.finish()
|
| 265 |
+
return
|
| 266 |
+
}
|
| 267 |
+
|
| 268 |
+
DispatchQueue.main.async {
|
| 269 |
+
self.state = .generating
|
| 270 |
+
self.tokensGenerated = 0
|
| 271 |
+
self.tokensPerSecond = 0
|
| 272 |
+
self.isGenerating = true
|
| 273 |
+
}
|
| 274 |
+
|
| 275 |
+
nonisolated(unsafe) let ctxForCancel = ctx
|
| 276 |
+
continuation.onTermination = { @Sendable _ in
|
| 277 |
+
flashmoe_cancel(ctxForCancel)
|
| 278 |
+
}
|
| 279 |
+
|
| 280 |
+
engineQueue.async { [weak self] in
|
| 281 |
+
let userDataPtr = Unmanaged.passRetained(
|
| 282 |
+
TokenCallbackContext(continuation: continuation, engine: self)
|
| 283 |
+
).toOpaque()
|
| 284 |
+
|
| 285 |
+
let result = flashmoe_generate_continuation(
|
| 286 |
+
ctx,
|
| 287 |
+
userMessage,
|
| 288 |
+
Int32(maxTokens),
|
| 289 |
+
{ tokenText, tokenId, tokensGenerated, tokensPerSecond, userData -> Int32 in
|
| 290 |
+
guard let userData else { return 1 }
|
| 291 |
+
let context = Unmanaged<TokenCallbackContext>.fromOpaque(userData)
|
| 292 |
+
.takeUnretainedValue()
|
| 293 |
+
|
| 294 |
+
guard let text = tokenText else { return 0 }
|
| 295 |
+
let token = GenerationToken(
|
| 296 |
+
text: String(cString: text),
|
| 297 |
+
tokenId: Int(tokenId),
|
| 298 |
+
tokensGenerated: Int(tokensGenerated),
|
| 299 |
+
tokensPerSecond: tokensPerSecond
|
| 300 |
+
)
|
| 301 |
+
|
| 302 |
+
if let engine = context.engine {
|
| 303 |
+
DispatchQueue.main.async {
|
| 304 |
+
engine.tokensGenerated = Int(tokensGenerated)
|
| 305 |
+
engine.tokensPerSecond = tokensPerSecond
|
| 306 |
+
}
|
| 307 |
+
}
|
| 308 |
+
|
| 309 |
+
context.continuation.yield(token)
|
| 310 |
+
return 0
|
| 311 |
+
},
|
| 312 |
+
userDataPtr
|
| 313 |
+
)
|
| 314 |
+
|
| 315 |
+
Unmanaged<TokenCallbackContext>.fromOpaque(userDataPtr).release()
|
| 316 |
+
|
| 317 |
+
// -2 = context full, signal via empty stream (caller handles reset)
|
| 318 |
+
if result == -2 {
|
| 319 |
+
DispatchQueue.main.async {
|
| 320 |
+
self?.state = .ready
|
| 321 |
+
self?.isGenerating = false
|
| 322 |
+
}
|
| 323 |
+
continuation.finish()
|
| 324 |
+
return
|
| 325 |
+
}
|
| 326 |
+
|
| 327 |
+
var stats = FlashMoEStats()
|
| 328 |
+
flashmoe_get_stats(ctx, &stats)
|
| 329 |
+
|
| 330 |
+
DispatchQueue.main.async {
|
| 331 |
+
self?.timeToFirstToken = stats.ttft_ms
|
| 332 |
+
self?.tokensPerSecond = stats.tokens_per_second
|
| 333 |
+
self?.tokensGenerated = Int(stats.tokens_generated)
|
| 334 |
+
self?.state = .ready
|
| 335 |
+
self?.isGenerating = false
|
| 336 |
+
}
|
| 337 |
+
|
| 338 |
+
continuation.finish()
|
| 339 |
+
}
|
| 340 |
+
}
|
| 341 |
+
}
|
| 342 |
+
|
| 343 |
+
/// Whether the engine has conversation state that can be continued
|
| 344 |
+
var canContinue: Bool {
|
| 345 |
+
guard let ctx = context else { return false }
|
| 346 |
+
return flashmoe_turn_count(ctx) > 0
|
| 347 |
+
}
|
| 348 |
+
|
| 349 |
+
/// Cancel an in-progress generation
|
| 350 |
+
func cancel() {
|
| 351 |
+
guard let ctx = context, isGenerating else { return }
|
| 352 |
+
flashmoe_cancel(ctx)
|
| 353 |
+
}
|
| 354 |
+
|
| 355 |
+
/// Reset conversation state (KV cache, attention state)
|
| 356 |
+
func reset() {
|
| 357 |
+
guard let ctx = context else { return }
|
| 358 |
+
engineQueue.async {
|
| 359 |
+
flashmoe_reset(ctx)
|
| 360 |
+
}
|
| 361 |
+
}
|
| 362 |
+
|
| 363 |
+
// MARK: - Runtime Config (no reload)
|
| 364 |
+
|
| 365 |
+
/// Change inference settings without reloading the model.
|
| 366 |
+
/// Sets C globals directly. Much faster than unload+reload.
|
| 367 |
+
func applyBenchmarkConfig(activeExpertsK: Int, cmdMerge: Bool,
|
| 368 |
+
fusedAttention: Bool, cacheIOSplit: Int,
|
| 369 |
+
fp16Accumulation: Bool) {
|
| 370 |
+
guard let ctx = context else { return }
|
| 371 |
+
flashmoe_set_runtime_config(ctx,
|
| 372 |
+
Int32(activeExpertsK),
|
| 373 |
+
cmdMerge ? 1 : 0,
|
| 374 |
+
fusedAttention ? 1 : 0,
|
| 375 |
+
Int32(cacheIOSplit),
|
| 376 |
+
fp16Accumulation ? 1 : 0)
|
| 377 |
+
}
|
| 378 |
+
|
| 379 |
+
// MARK: - Model Validation
|
| 380 |
+
|
| 381 |
+
/// Check if a model directory contains a valid Flash-MoE model
|
| 382 |
+
static func validateModel(at path: String) -> Bool {
|
| 383 |
+
return flashmoe_validate_model(path) == 0
|
| 384 |
+
}
|
| 385 |
+
}
|
| 386 |
+
|
| 387 |
+
// MARK: - Helper Types
|
| 388 |
+
|
| 389 |
+
/// Bridging class to pass Swift state through C void* callback
|
| 390 |
+
private final class TokenCallbackContext {
|
| 391 |
+
let continuation: AsyncStream<GenerationToken>.Continuation
|
| 392 |
+
weak var engine: FlashMoEEngine?
|
| 393 |
+
|
| 394 |
+
init(continuation: AsyncStream<GenerationToken>.Continuation, engine: FlashMoEEngine?) {
|
| 395 |
+
self.continuation = continuation
|
| 396 |
+
self.engine = engine
|
| 397 |
+
}
|
| 398 |
+
}
|
| 399 |
+
|
| 400 |
+
/// Errors from the Flash-MoE engine
|
| 401 |
+
enum FlashMoEError: LocalizedError {
|
| 402 |
+
case busy
|
| 403 |
+
case engineDestroyed
|
| 404 |
+
case initFailed
|
| 405 |
+
case loadFailed(String)
|
| 406 |
+
case generateFailed(String)
|
| 407 |
+
case notLoaded
|
| 408 |
+
|
| 409 |
+
var errorDescription: String? {
|
| 410 |
+
switch self {
|
| 411 |
+
case .busy: return "Engine is busy"
|
| 412 |
+
case .engineDestroyed: return "Engine was destroyed"
|
| 413 |
+
case .initFailed: return "Failed to initialize engine"
|
| 414 |
+
case .loadFailed(let msg): return "Failed to load model: \(msg)"
|
| 415 |
+
case .generateFailed(let msg): return "Generation failed: \(msg)"
|
| 416 |
+
case .notLoaded: return "No model loaded"
|
| 417 |
+
}
|
| 418 |
+
}
|
| 419 |
+
}
|
FlashMoE-iOS/FlashMoE.xcodeproj/project.pbxproj
ADDED
|
@@ -0,0 +1,516 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// !$*UTF8*$!
|
| 2 |
+
{
|
| 3 |
+
archiveVersion = 1;
|
| 4 |
+
classes = {
|
| 5 |
+
};
|
| 6 |
+
objectVersion = 77;
|
| 7 |
+
objects = {
|
| 8 |
+
|
| 9 |
+
/* Begin PBXBuildFile section */
|
| 10 |
+
FF1111111111111111111111 /* BenchmarkView.swift in Sources */ = {isa = PBXBuildFile; fileRef = FF2222222222222222222222 /* BenchmarkView.swift */; };
|
| 11 |
+
0BA5B28A95E8C0616CCA6A3D /* FlashMoEApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8EAF41B16BD4314D3642DC53 /* FlashMoEApp.swift */; };
|
| 12 |
+
15BBE610777A6EF58279A4E3 /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2AEBDCE15020A40A8451CB5A /* ContentView.swift */; };
|
| 13 |
+
23E305445AFD3D5FB1A5E5FC /* ModelListView.swift in Sources */ = {isa = PBXBuildFile; fileRef = DDD107A34A4DAEABC7E9ECFB /* ModelListView.swift */; };
|
| 14 |
+
368893C82D55ADC634FBAD8B /* Metal.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 8A8F6936F5C735B8856277A2 /* Metal.framework */; };
|
| 15 |
+
5F7FD44C4FF2E97505EDAB58 /* FlashMoEBridge.swift in Sources */ = {isa = PBXBuildFile; fileRef = 749C35B05B7C71B73796EEE6 /* FlashMoEBridge.swift */; };
|
| 16 |
+
86FB8D83EC8F2AA092081685 /* ChatView.swift in Sources */ = {isa = PBXBuildFile; fileRef = BAC1A43E0715A9CDE325FFC0 /* ChatView.swift */; };
|
| 17 |
+
A40E8D436A3EEDB8B04CDC44 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = DEE08D3053B24D1E206627F9 /* Assets.xcassets */; };
|
| 18 |
+
AA1111111111111111111111 /* ModelCatalog.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA2222222222222222222222 /* ModelCatalog.swift */; };
|
| 19 |
+
BB1111111111111111111111 /* DownloadManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = BB2222222222222222222222 /* DownloadManager.swift */; };
|
| 20 |
+
CC1111111111111111111111 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = CC2222222222222222222222 /* AppDelegate.swift */; };
|
| 21 |
+
D2F140EBFBCC62B37BDBCB7F /* Foundation.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 3321AD5929C6B6D4FAD446C8 /* Foundation.framework */; };
|
| 22 |
+
D8242F5999552BC228A650CE /* shaders.metal in Sources */ = {isa = PBXBuildFile; fileRef = EEC822B6F9CE0BABA49E989B /* shaders.metal */; };
|
| 23 |
+
DD1111111111111111111111 /* ModelDownloadRow.swift in Sources */ = {isa = PBXBuildFile; fileRef = DD2222222222222222222222 /* ModelDownloadRow.swift */; };
|
| 24 |
+
E73BB005AD9F878023B7F75E /* Accelerate.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 51FB34EBC20226BA8269877A /* Accelerate.framework */; };
|
| 25 |
+
EE1111111111111111111111 /* ProfilerView.swift in Sources */ = {isa = PBXBuildFile; fileRef = EE2222222222222222222222 /* ProfilerView.swift */; };
|
| 26 |
+
F00087500AE558979B35EE06 /* FlashMoEEngine.m in Sources */ = {isa = PBXBuildFile; fileRef = 50DF36AC37579556033D375F /* FlashMoEEngine.m */; };
|
| 27 |
+
/* End PBXBuildFile section */
|
| 28 |
+
|
| 29 |
+
/* Begin PBXFileReference section */
|
| 30 |
+
03B08EB0D0CE594165AB1C14 /* FlashMoE-Bridging-Header.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = "FlashMoE-Bridging-Header.h"; sourceTree = "<group>"; };
|
| 31 |
+
FF2222222222222222222222 /* BenchmarkView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BenchmarkView.swift; sourceTree = "<group>"; };
|
| 32 |
+
2AEBDCE15020A40A8451CB5A /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = "<group>"; };
|
| 33 |
+
3321AD5929C6B6D4FAD446C8 /* Foundation.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Foundation.framework; path = System/Library/Frameworks/Foundation.framework; sourceTree = SDKROOT; };
|
| 34 |
+
3404570FACF1297B40EF6369 /* FlashMoEEngine.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = FlashMoEEngine.h; sourceTree = "<group>"; };
|
| 35 |
+
50DF36AC37579556033D375F /* FlashMoEEngine.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = FlashMoEEngine.m; sourceTree = "<group>"; };
|
| 36 |
+
51FB34EBC20226BA8269877A /* Accelerate.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Accelerate.framework; path = System/Library/Frameworks/Accelerate.framework; sourceTree = SDKROOT; };
|
| 37 |
+
575B85659731B0C7F170E59E /* tokenizer.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = tokenizer.h; sourceTree = "<group>"; };
|
| 38 |
+
64E904A15D9438F2D9B365F5 /* FlashMoE.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = FlashMoE.app; sourceTree = BUILT_PRODUCTS_DIR; };
|
| 39 |
+
749C35B05B7C71B73796EEE6 /* FlashMoEBridge.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FlashMoEBridge.swift; sourceTree = "<group>"; };
|
| 40 |
+
82BDA279F2267A8FD313D3C3 /* FlashMoE.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = FlashMoE.entitlements; sourceTree = "<group>"; };
|
| 41 |
+
8A8F6936F5C735B8856277A2 /* Metal.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Metal.framework; path = System/Library/Frameworks/Metal.framework; sourceTree = SDKROOT; };
|
| 42 |
+
8EAF41B16BD4314D3642DC53 /* FlashMoEApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FlashMoEApp.swift; sourceTree = "<group>"; };
|
| 43 |
+
AA2222222222222222222222 /* ModelCatalog.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelCatalog.swift; sourceTree = "<group>"; };
|
| 44 |
+
BAC1A43E0715A9CDE325FFC0 /* ChatView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ChatView.swift; sourceTree = "<group>"; };
|
| 45 |
+
BB2222222222222222222222 /* DownloadManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = DownloadManager.swift; sourceTree = "<group>"; };
|
| 46 |
+
CC2222222222222222222222 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = "<group>"; };
|
| 47 |
+
DD2222222222222222222222 /* ModelDownloadRow.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelDownloadRow.swift; sourceTree = "<group>"; };
|
| 48 |
+
DDD107A34A4DAEABC7E9ECFB /* ModelListView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelListView.swift; sourceTree = "<group>"; };
|
| 49 |
+
DEE08D3053B24D1E206627F9 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
|
| 50 |
+
EE2222222222222222222222 /* ProfilerView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ProfilerView.swift; sourceTree = "<group>"; };
|
| 51 |
+
EEC822B6F9CE0BABA49E989B /* shaders.metal */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.metal; path = shaders.metal; sourceTree = "<group>"; };
|
| 52 |
+
/* End PBXFileReference section */
|
| 53 |
+
|
| 54 |
+
/* Begin PBXFrameworksBuildPhase section */
|
| 55 |
+
FD2BB59A0A54FF99587D143D /* Frameworks */ = {
|
| 56 |
+
isa = PBXFrameworksBuildPhase;
|
| 57 |
+
buildActionMask = 2147483647;
|
| 58 |
+
files = (
|
| 59 |
+
368893C82D55ADC634FBAD8B /* Metal.framework in Frameworks */,
|
| 60 |
+
D2F140EBFBCC62B37BDBCB7F /* Foundation.framework in Frameworks */,
|
| 61 |
+
E73BB005AD9F878023B7F75E /* Accelerate.framework in Frameworks */,
|
| 62 |
+
);
|
| 63 |
+
runOnlyForDeploymentPostprocessing = 0;
|
| 64 |
+
};
|
| 65 |
+
/* End PBXFrameworksBuildPhase section */
|
| 66 |
+
|
| 67 |
+
/* Begin PBXGroup section */
|
| 68 |
+
18638B2F96B8C1EB7D55C874 /* Bridge */ = {
|
| 69 |
+
isa = PBXGroup;
|
| 70 |
+
children = (
|
| 71 |
+
03B08EB0D0CE594165AB1C14 /* FlashMoE-Bridging-Header.h */,
|
| 72 |
+
749C35B05B7C71B73796EEE6 /* FlashMoEBridge.swift */,
|
| 73 |
+
);
|
| 74 |
+
path = Bridge;
|
| 75 |
+
sourceTree = "<group>";
|
| 76 |
+
};
|
| 77 |
+
1BE036041A7213E9AE43E96C /* Products */ = {
|
| 78 |
+
isa = PBXGroup;
|
| 79 |
+
children = (
|
| 80 |
+
64E904A15D9438F2D9B365F5 /* FlashMoE.app */,
|
| 81 |
+
);
|
| 82 |
+
name = Products;
|
| 83 |
+
sourceTree = "<group>";
|
| 84 |
+
};
|
| 85 |
+
42B268D8ECA1DA8E59B4F599 /* ViewModels */ = {
|
| 86 |
+
isa = PBXGroup;
|
| 87 |
+
children = (
|
| 88 |
+
);
|
| 89 |
+
path = ViewModels;
|
| 90 |
+
sourceTree = "<group>";
|
| 91 |
+
};
|
| 92 |
+
436D67F159A3C0FA4F05317F /* metal_infer */ = {
|
| 93 |
+
isa = PBXGroup;
|
| 94 |
+
children = (
|
| 95 |
+
EEC822B6F9CE0BABA49E989B /* shaders.metal */,
|
| 96 |
+
575B85659731B0C7F170E59E /* tokenizer.h */,
|
| 97 |
+
);
|
| 98 |
+
name = metal_infer;
|
| 99 |
+
path = ../metal_infer;
|
| 100 |
+
sourceTree = "<group>";
|
| 101 |
+
};
|
| 102 |
+
5A512650C4615EA193374D2F /* Resources */ = {
|
| 103 |
+
isa = PBXGroup;
|
| 104 |
+
children = (
|
| 105 |
+
DEE08D3053B24D1E206627F9 /* Assets.xcassets */,
|
| 106 |
+
82BDA279F2267A8FD313D3C3 /* FlashMoE.entitlements */,
|
| 107 |
+
);
|
| 108 |
+
path = Resources;
|
| 109 |
+
sourceTree = "<group>";
|
| 110 |
+
};
|
| 111 |
+
7066FB671F127A31C9768B99 /* Frameworks */ = {
|
| 112 |
+
isa = PBXGroup;
|
| 113 |
+
children = (
|
| 114 |
+
51FB34EBC20226BA8269877A /* Accelerate.framework */,
|
| 115 |
+
3321AD5929C6B6D4FAD446C8 /* Foundation.framework */,
|
| 116 |
+
8A8F6936F5C735B8856277A2 /* Metal.framework */,
|
| 117 |
+
);
|
| 118 |
+
name = Frameworks;
|
| 119 |
+
sourceTree = "<group>";
|
| 120 |
+
};
|
| 121 |
+
7DC3D96AAD773F68496A645B /* App */ = {
|
| 122 |
+
isa = PBXGroup;
|
| 123 |
+
children = (
|
| 124 |
+
CC2222222222222222222222 /* AppDelegate.swift */,
|
| 125 |
+
8EAF41B16BD4314D3642DC53 /* FlashMoEApp.swift */,
|
| 126 |
+
);
|
| 127 |
+
path = App;
|
| 128 |
+
sourceTree = "<group>";
|
| 129 |
+
};
|
| 130 |
+
9072059C570A7ED56A540CBE = {
|
| 131 |
+
isa = PBXGroup;
|
| 132 |
+
children = (
|
| 133 |
+
7DC3D96AAD773F68496A645B /* App */,
|
| 134 |
+
18638B2F96B8C1EB7D55C874 /* Bridge */,
|
| 135 |
+
EDA325EB4F586262787DC634 /* FlashMoEEngine */,
|
| 136 |
+
436D67F159A3C0FA4F05317F /* metal_infer */,
|
| 137 |
+
AA3333333333333333333333 /* Models */,
|
| 138 |
+
5A512650C4615EA193374D2F /* Resources */,
|
| 139 |
+
BB3333333333333333333333 /* Services */,
|
| 140 |
+
42B268D8ECA1DA8E59B4F599 /* ViewModels */,
|
| 141 |
+
993DF219F89AA4229E9A43D2 /* Views */,
|
| 142 |
+
7066FB671F127A31C9768B99 /* Frameworks */,
|
| 143 |
+
1BE036041A7213E9AE43E96C /* Products */,
|
| 144 |
+
);
|
| 145 |
+
sourceTree = "<group>";
|
| 146 |
+
};
|
| 147 |
+
993DF219F89AA4229E9A43D2 /* Views */ = {
|
| 148 |
+
isa = PBXGroup;
|
| 149 |
+
children = (
|
| 150 |
+
FF2222222222222222222222 /* BenchmarkView.swift */,
|
| 151 |
+
BAC1A43E0715A9CDE325FFC0 /* ChatView.swift */,
|
| 152 |
+
2AEBDCE15020A40A8451CB5A /* ContentView.swift */,
|
| 153 |
+
DD2222222222222222222222 /* ModelDownloadRow.swift */,
|
| 154 |
+
DDD107A34A4DAEABC7E9ECFB /* ModelListView.swift */,
|
| 155 |
+
EE2222222222222222222222 /* ProfilerView.swift */,
|
| 156 |
+
);
|
| 157 |
+
path = Views;
|
| 158 |
+
sourceTree = "<group>";
|
| 159 |
+
};
|
| 160 |
+
AA3333333333333333333333 /* Models */ = {
|
| 161 |
+
isa = PBXGroup;
|
| 162 |
+
children = (
|
| 163 |
+
AA2222222222222222222222 /* ModelCatalog.swift */,
|
| 164 |
+
);
|
| 165 |
+
path = Models;
|
| 166 |
+
sourceTree = "<group>";
|
| 167 |
+
};
|
| 168 |
+
BB3333333333333333333333 /* Services */ = {
|
| 169 |
+
isa = PBXGroup;
|
| 170 |
+
children = (
|
| 171 |
+
BB2222222222222222222222 /* DownloadManager.swift */,
|
| 172 |
+
);
|
| 173 |
+
path = Services;
|
| 174 |
+
sourceTree = "<group>";
|
| 175 |
+
};
|
| 176 |
+
EDA325EB4F586262787DC634 /* FlashMoEEngine */ = {
|
| 177 |
+
isa = PBXGroup;
|
| 178 |
+
children = (
|
| 179 |
+
3404570FACF1297B40EF6369 /* FlashMoEEngine.h */,
|
| 180 |
+
50DF36AC37579556033D375F /* FlashMoEEngine.m */,
|
| 181 |
+
);
|
| 182 |
+
path = FlashMoEEngine;
|
| 183 |
+
sourceTree = "<group>";
|
| 184 |
+
};
|
| 185 |
+
/* End PBXGroup section */
|
| 186 |
+
|
| 187 |
+
/* Begin PBXNativeTarget section */
|
| 188 |
+
F64C8CE03661205C4C06D1FF /* FlashMoE */ = {
|
| 189 |
+
isa = PBXNativeTarget;
|
| 190 |
+
buildConfigurationList = A992692CC1A07BC40EFA0D3E /* Build configuration list for PBXNativeTarget "FlashMoE" */;
|
| 191 |
+
buildPhases = (
|
| 192 |
+
D57FD46C656820096AA63872 /* Sources */,
|
| 193 |
+
E55762DD03B859815792AC85 /* Resources */,
|
| 194 |
+
FD2BB59A0A54FF99587D143D /* Frameworks */,
|
| 195 |
+
);
|
| 196 |
+
buildRules = (
|
| 197 |
+
);
|
| 198 |
+
dependencies = (
|
| 199 |
+
);
|
| 200 |
+
name = FlashMoE;
|
| 201 |
+
packageProductDependencies = (
|
| 202 |
+
);
|
| 203 |
+
productName = FlashMoE;
|
| 204 |
+
productReference = 64E904A15D9438F2D9B365F5 /* FlashMoE.app */;
|
| 205 |
+
productType = "com.apple.product-type.application";
|
| 206 |
+
};
|
| 207 |
+
/* End PBXNativeTarget section */
|
| 208 |
+
|
| 209 |
+
/* Begin PBXProject section */
|
| 210 |
+
D6A8CACD08FF01EA2A8633B6 /* Project object */ = {
|
| 211 |
+
isa = PBXProject;
|
| 212 |
+
attributes = {
|
| 213 |
+
BuildIndependentTargetsInParallel = YES;
|
| 214 |
+
LastUpgradeCheck = 1600;
|
| 215 |
+
};
|
| 216 |
+
buildConfigurationList = 716F6ADA7E507FC5C74161A0 /* Build configuration list for PBXProject "FlashMoE" */;
|
| 217 |
+
developmentRegion = en;
|
| 218 |
+
hasScannedForEncodings = 0;
|
| 219 |
+
knownRegions = (
|
| 220 |
+
Base,
|
| 221 |
+
en,
|
| 222 |
+
);
|
| 223 |
+
mainGroup = 9072059C570A7ED56A540CBE;
|
| 224 |
+
minimizedProjectReferenceProxies = 1;
|
| 225 |
+
preferredProjectObjectVersion = 77;
|
| 226 |
+
productRefGroup = 1BE036041A7213E9AE43E96C /* Products */;
|
| 227 |
+
projectDirPath = "";
|
| 228 |
+
projectRoot = "";
|
| 229 |
+
targets = (
|
| 230 |
+
F64C8CE03661205C4C06D1FF /* FlashMoE */,
|
| 231 |
+
);
|
| 232 |
+
};
|
| 233 |
+
/* End PBXProject section */
|
| 234 |
+
|
| 235 |
+
/* Begin PBXResourcesBuildPhase section */
|
| 236 |
+
E55762DD03B859815792AC85 /* Resources */ = {
|
| 237 |
+
isa = PBXResourcesBuildPhase;
|
| 238 |
+
buildActionMask = 2147483647;
|
| 239 |
+
files = (
|
| 240 |
+
A40E8D436A3EEDB8B04CDC44 /* Assets.xcassets in Resources */,
|
| 241 |
+
);
|
| 242 |
+
runOnlyForDeploymentPostprocessing = 0;
|
| 243 |
+
};
|
| 244 |
+
/* End PBXResourcesBuildPhase section */
|
| 245 |
+
|
| 246 |
+
/* Begin PBXSourcesBuildPhase section */
|
| 247 |
+
D57FD46C656820096AA63872 /* Sources */ = {
|
| 248 |
+
isa = PBXSourcesBuildPhase;
|
| 249 |
+
buildActionMask = 2147483647;
|
| 250 |
+
files = (
|
| 251 |
+
D8242F5999552BC228A650CE /* shaders.metal in Sources */,
|
| 252 |
+
CC1111111111111111111111 /* AppDelegate.swift in Sources */,
|
| 253 |
+
FF1111111111111111111111 /* BenchmarkView.swift in Sources */,
|
| 254 |
+
86FB8D83EC8F2AA092081685 /* ChatView.swift in Sources */,
|
| 255 |
+
15BBE610777A6EF58279A4E3 /* ContentView.swift in Sources */,
|
| 256 |
+
BB1111111111111111111111 /* DownloadManager.swift in Sources */,
|
| 257 |
+
0BA5B28A95E8C0616CCA6A3D /* FlashMoEApp.swift in Sources */,
|
| 258 |
+
5F7FD44C4FF2E97505EDAB58 /* FlashMoEBridge.swift in Sources */,
|
| 259 |
+
F00087500AE558979B35EE06 /* FlashMoEEngine.m in Sources */,
|
| 260 |
+
AA1111111111111111111111 /* ModelCatalog.swift in Sources */,
|
| 261 |
+
DD1111111111111111111111 /* ModelDownloadRow.swift in Sources */,
|
| 262 |
+
23E305445AFD3D5FB1A5E5FC /* ModelListView.swift in Sources */,
|
| 263 |
+
EE1111111111111111111111 /* ProfilerView.swift in Sources */,
|
| 264 |
+
);
|
| 265 |
+
runOnlyForDeploymentPostprocessing = 0;
|
| 266 |
+
};
|
| 267 |
+
/* End PBXSourcesBuildPhase section */
|
| 268 |
+
|
| 269 |
+
/* Begin XCBuildConfiguration section */
|
| 270 |
+
20C11034B98051ECF5657F59 /* Release */ = {
|
| 271 |
+
isa = XCBuildConfiguration;
|
| 272 |
+
buildSettings = {
|
| 273 |
+
ALWAYS_SEARCH_USER_PATHS = NO;
|
| 274 |
+
CLANG_ANALYZER_NONNULL = YES;
|
| 275 |
+
CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
|
| 276 |
+
CLANG_CXX_LANGUAGE_STANDARD = "gnu++14";
|
| 277 |
+
CLANG_CXX_LIBRARY = "libc++";
|
| 278 |
+
CLANG_ENABLE_MODULES = YES;
|
| 279 |
+
CLANG_ENABLE_OBJC_ARC = YES;
|
| 280 |
+
CLANG_ENABLE_OBJC_WEAK = YES;
|
| 281 |
+
CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
|
| 282 |
+
CLANG_WARN_BOOL_CONVERSION = YES;
|
| 283 |
+
CLANG_WARN_COMMA = YES;
|
| 284 |
+
CLANG_WARN_CONSTANT_CONVERSION = YES;
|
| 285 |
+
CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
|
| 286 |
+
CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
|
| 287 |
+
CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
|
| 288 |
+
CLANG_WARN_EMPTY_BODY = YES;
|
| 289 |
+
CLANG_WARN_ENUM_CONVERSION = YES;
|
| 290 |
+
CLANG_WARN_INFINITE_RECURSION = YES;
|
| 291 |
+
CLANG_WARN_INT_CONVERSION = YES;
|
| 292 |
+
CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
|
| 293 |
+
CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
|
| 294 |
+
CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
|
| 295 |
+
CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
|
| 296 |
+
CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
|
| 297 |
+
CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
|
| 298 |
+
CLANG_WARN_STRICT_PROTOTYPES = YES;
|
| 299 |
+
CLANG_WARN_SUSPICIOUS_MOVE = YES;
|
| 300 |
+
CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
|
| 301 |
+
CLANG_WARN_UNREACHABLE_CODE = YES;
|
| 302 |
+
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
|
| 303 |
+
COPY_PHASE_STRIP = NO;
|
| 304 |
+
DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
|
| 305 |
+
ENABLE_NS_ASSERTIONS = NO;
|
| 306 |
+
ENABLE_STRICT_OBJC_MSGSEND = YES;
|
| 307 |
+
GCC_C_LANGUAGE_STANDARD = gnu11;
|
| 308 |
+
GCC_NO_COMMON_BLOCKS = YES;
|
| 309 |
+
GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
|
| 310 |
+
GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
|
| 311 |
+
GCC_WARN_UNDECLARED_SELECTOR = YES;
|
| 312 |
+
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
|
| 313 |
+
GCC_WARN_UNUSED_FUNCTION = YES;
|
| 314 |
+
GCC_WARN_UNUSED_VARIABLE = YES;
|
| 315 |
+
INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace = YES;
|
| 316 |
+
INFOPLIST_KEY_UIFileSharingEnabled = YES;
|
| 317 |
+
IPHONEOS_DEPLOYMENT_TARGET = 18.0;
|
| 318 |
+
MACOSX_DEPLOYMENT_TARGET = 15.0;
|
| 319 |
+
MTL_ENABLE_DEBUG_INFO = NO;
|
| 320 |
+
MTL_FAST_MATH = YES;
|
| 321 |
+
PRODUCT_NAME = "$(TARGET_NAME)";
|
| 322 |
+
SDKROOT = auto;
|
| 323 |
+
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
|
| 324 |
+
SWIFT_COMPILATION_MODE = wholemodule;
|
| 325 |
+
SWIFT_OPTIMIZATION_LEVEL = "-O";
|
| 326 |
+
SWIFT_VERSION = 6.0;
|
| 327 |
+
TARGETED_DEVICE_FAMILY = "1,2";
|
| 328 |
+
};
|
| 329 |
+
name = Release;
|
| 330 |
+
};
|
| 331 |
+
8C87D5217C25485D83AEE132 /* Debug */ = {
|
| 332 |
+
isa = XCBuildConfiguration;
|
| 333 |
+
buildSettings = {
|
| 334 |
+
ALWAYS_SEARCH_USER_PATHS = NO;
|
| 335 |
+
CLANG_ANALYZER_NONNULL = YES;
|
| 336 |
+
CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
|
| 337 |
+
CLANG_CXX_LANGUAGE_STANDARD = "gnu++14";
|
| 338 |
+
CLANG_CXX_LIBRARY = "libc++";
|
| 339 |
+
CLANG_ENABLE_MODULES = YES;
|
| 340 |
+
CLANG_ENABLE_OBJC_ARC = YES;
|
| 341 |
+
CLANG_ENABLE_OBJC_WEAK = YES;
|
| 342 |
+
CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
|
| 343 |
+
CLANG_WARN_BOOL_CONVERSION = YES;
|
| 344 |
+
CLANG_WARN_COMMA = YES;
|
| 345 |
+
CLANG_WARN_CONSTANT_CONVERSION = YES;
|
| 346 |
+
CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
|
| 347 |
+
CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
|
| 348 |
+
CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
|
| 349 |
+
CLANG_WARN_EMPTY_BODY = YES;
|
| 350 |
+
CLANG_WARN_ENUM_CONVERSION = YES;
|
| 351 |
+
CLANG_WARN_INFINITE_RECURSION = YES;
|
| 352 |
+
CLANG_WARN_INT_CONVERSION = YES;
|
| 353 |
+
CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
|
| 354 |
+
CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
|
| 355 |
+
CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
|
| 356 |
+
CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
|
| 357 |
+
CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
|
| 358 |
+
CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
|
| 359 |
+
CLANG_WARN_STRICT_PROTOTYPES = YES;
|
| 360 |
+
CLANG_WARN_SUSPICIOUS_MOVE = YES;
|
| 361 |
+
CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
|
| 362 |
+
CLANG_WARN_UNREACHABLE_CODE = YES;
|
| 363 |
+
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
|
| 364 |
+
COPY_PHASE_STRIP = NO;
|
| 365 |
+
DEBUG_INFORMATION_FORMAT = dwarf;
|
| 366 |
+
ENABLE_STRICT_OBJC_MSGSEND = YES;
|
| 367 |
+
ENABLE_TESTABILITY = YES;
|
| 368 |
+
GCC_C_LANGUAGE_STANDARD = gnu11;
|
| 369 |
+
GCC_DYNAMIC_NO_PIC = NO;
|
| 370 |
+
GCC_NO_COMMON_BLOCKS = YES;
|
| 371 |
+
GCC_OPTIMIZATION_LEVEL = 0;
|
| 372 |
+
GCC_PREPROCESSOR_DEFINITIONS = (
|
| 373 |
+
"$(inherited)",
|
| 374 |
+
"DEBUG=1",
|
| 375 |
+
);
|
| 376 |
+
GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
|
| 377 |
+
GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
|
| 378 |
+
GCC_WARN_UNDECLARED_SELECTOR = YES;
|
| 379 |
+
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
|
| 380 |
+
GCC_WARN_UNUSED_FUNCTION = YES;
|
| 381 |
+
GCC_WARN_UNUSED_VARIABLE = YES;
|
| 382 |
+
INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace = YES;
|
| 383 |
+
INFOPLIST_KEY_UIFileSharingEnabled = YES;
|
| 384 |
+
IPHONEOS_DEPLOYMENT_TARGET = 18.0;
|
| 385 |
+
MACOSX_DEPLOYMENT_TARGET = 15.0;
|
| 386 |
+
MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
|
| 387 |
+
MTL_FAST_MATH = YES;
|
| 388 |
+
ONLY_ACTIVE_ARCH = YES;
|
| 389 |
+
PRODUCT_NAME = "$(TARGET_NAME)";
|
| 390 |
+
SDKROOT = auto;
|
| 391 |
+
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
|
| 392 |
+
SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG;
|
| 393 |
+
SWIFT_OPTIMIZATION_LEVEL = "-Onone";
|
| 394 |
+
SWIFT_VERSION = 6.0;
|
| 395 |
+
TARGETED_DEVICE_FAMILY = "1,2";
|
| 396 |
+
};
|
| 397 |
+
name = Debug;
|
| 398 |
+
};
|
| 399 |
+
A93B546F9091A305E7C156B1 /* Release */ = {
|
| 400 |
+
isa = XCBuildConfiguration;
|
| 401 |
+
buildSettings = {
|
| 402 |
+
ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
|
| 403 |
+
CLANG_ENABLE_OBJC_ARC = YES;
|
| 404 |
+
CODE_SIGN_ENTITLEMENTS = Resources/FlashMoE.entitlements;
|
| 405 |
+
CODE_SIGN_IDENTITY = "iPhone Developer";
|
| 406 |
+
"CODE_SIGN_IDENTITY[sdk=macosx*]" = "-";
|
| 407 |
+
DEVELOPMENT_TEAM = YMCRS53QKE;
|
| 408 |
+
GCC_OPTIMIZATION_LEVEL = 2;
|
| 409 |
+
GCC_WARN_INHIBIT_ALL_WARNINGS = NO;
|
| 410 |
+
GENERATE_INFOPLIST_FILE = YES;
|
| 411 |
+
HEADER_SEARCH_PATHS = (
|
| 412 |
+
"$(SRCROOT)/FlashMoEEngine",
|
| 413 |
+
"$(SRCROOT)/../metal_infer",
|
| 414 |
+
);
|
| 415 |
+
INFOPLIST_KEY_CFBundleDisplayName = "Flash-MoE";
|
| 416 |
+
INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace = YES;
|
| 417 |
+
INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES;
|
| 418 |
+
INFOPLIST_KEY_UILaunchScreen_Generation = YES;
|
| 419 |
+
INFOPLIST_KEY_UISupportedInterfaceOrientations = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
|
| 420 |
+
IPHONEOS_DEPLOYMENT_TARGET = 26.0;
|
| 421 |
+
LD_RUNPATH_SEARCH_PATHS = (
|
| 422 |
+
"$(inherited)",
|
| 423 |
+
"@executable_path/Frameworks",
|
| 424 |
+
"@executable_path/../Frameworks",
|
| 425 |
+
);
|
| 426 |
+
MACOSX_DEPLOYMENT_TARGET = 15.0;
|
| 427 |
+
OTHER_CFLAGS = (
|
| 428 |
+
"-DCHAT_MODE=1",
|
| 429 |
+
"-DACCELERATE_NEW_LAPACK",
|
| 430 |
+
);
|
| 431 |
+
OTHER_LDFLAGS = "-lcompression";
|
| 432 |
+
PRODUCT_BUNDLE_IDENTIFIER = com.flashmoe.app;
|
| 433 |
+
PRODUCT_NAME = FlashMoE;
|
| 434 |
+
SDKROOT = auto;
|
| 435 |
+
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
|
| 436 |
+
SUPPORTS_MACCATALYST = NO;
|
| 437 |
+
SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD = NO;
|
| 438 |
+
SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD = NO;
|
| 439 |
+
SWIFT_OBJC_BRIDGING_HEADER = "Bridge/FlashMoE-Bridging-Header.h";
|
| 440 |
+
TARGETED_DEVICE_FAMILY = "1,2";
|
| 441 |
+
};
|
| 442 |
+
name = Release;
|
| 443 |
+
};
|
| 444 |
+
F06D0A6E051D510D77122F00 /* Debug */ = {
|
| 445 |
+
isa = XCBuildConfiguration;
|
| 446 |
+
buildSettings = {
|
| 447 |
+
ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
|
| 448 |
+
CLANG_ENABLE_OBJC_ARC = YES;
|
| 449 |
+
CODE_SIGN_ENTITLEMENTS = Resources/FlashMoE.entitlements;
|
| 450 |
+
CODE_SIGN_IDENTITY = "iPhone Developer";
|
| 451 |
+
"CODE_SIGN_IDENTITY[sdk=macosx*]" = "-";
|
| 452 |
+
DEVELOPMENT_TEAM = YMCRS53QKE;
|
| 453 |
+
ENABLE_APP_SANDBOX = YES;
|
| 454 |
+
ENABLE_HARDENED_RUNTIME = YES;
|
| 455 |
+
ENABLE_OUTGOING_NETWORK_CONNECTIONS = YES;
|
| 456 |
+
GCC_OPTIMIZATION_LEVEL = 2;
|
| 457 |
+
GCC_WARN_INHIBIT_ALL_WARNINGS = NO;
|
| 458 |
+
GENERATE_INFOPLIST_FILE = YES;
|
| 459 |
+
HEADER_SEARCH_PATHS = (
|
| 460 |
+
"$(SRCROOT)/FlashMoEEngine",
|
| 461 |
+
"$(SRCROOT)/../metal_infer",
|
| 462 |
+
);
|
| 463 |
+
INFOPLIST_KEY_CFBundleDisplayName = "Flash-MoE";
|
| 464 |
+
INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace = YES;
|
| 465 |
+
INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES;
|
| 466 |
+
INFOPLIST_KEY_UILaunchScreen_Generation = YES;
|
| 467 |
+
INFOPLIST_KEY_UISupportedInterfaceOrientations = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
|
| 468 |
+
IPHONEOS_DEPLOYMENT_TARGET = 26.0;
|
| 469 |
+
LD_RUNPATH_SEARCH_PATHS = (
|
| 470 |
+
"$(inherited)",
|
| 471 |
+
"@executable_path/Frameworks",
|
| 472 |
+
"@executable_path/../Frameworks",
|
| 473 |
+
);
|
| 474 |
+
MACOSX_DEPLOYMENT_TARGET = 15.0;
|
| 475 |
+
OTHER_CFLAGS = (
|
| 476 |
+
"-DCHAT_MODE=1",
|
| 477 |
+
"-DACCELERATE_NEW_LAPACK",
|
| 478 |
+
);
|
| 479 |
+
OTHER_LDFLAGS = "-lcompression";
|
| 480 |
+
PRODUCT_BUNDLE_IDENTIFIER = com.flashmoe.app;
|
| 481 |
+
PRODUCT_NAME = FlashMoE;
|
| 482 |
+
SDKROOT = auto;
|
| 483 |
+
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
|
| 484 |
+
SUPPORTS_MACCATALYST = NO;
|
| 485 |
+
SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD = NO;
|
| 486 |
+
SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD = NO;
|
| 487 |
+
SWIFT_OBJC_BRIDGING_HEADER = "Bridge/FlashMoE-Bridging-Header.h";
|
| 488 |
+
TARGETED_DEVICE_FAMILY = "1,2";
|
| 489 |
+
};
|
| 490 |
+
name = Debug;
|
| 491 |
+
};
|
| 492 |
+
/* End XCBuildConfiguration section */
|
| 493 |
+
|
| 494 |
+
/* Begin XCConfigurationList section */
|
| 495 |
+
716F6ADA7E507FC5C74161A0 /* Build configuration list for PBXProject "FlashMoE" */ = {
|
| 496 |
+
isa = XCConfigurationList;
|
| 497 |
+
buildConfigurations = (
|
| 498 |
+
8C87D5217C25485D83AEE132 /* Debug */,
|
| 499 |
+
20C11034B98051ECF5657F59 /* Release */,
|
| 500 |
+
);
|
| 501 |
+
defaultConfigurationIsVisible = 0;
|
| 502 |
+
defaultConfigurationName = Debug;
|
| 503 |
+
};
|
| 504 |
+
A992692CC1A07BC40EFA0D3E /* Build configuration list for PBXNativeTarget "FlashMoE" */ = {
|
| 505 |
+
isa = XCConfigurationList;
|
| 506 |
+
buildConfigurations = (
|
| 507 |
+
F06D0A6E051D510D77122F00 /* Debug */,
|
| 508 |
+
A93B546F9091A305E7C156B1 /* Release */,
|
| 509 |
+
);
|
| 510 |
+
defaultConfigurationIsVisible = 0;
|
| 511 |
+
defaultConfigurationName = Debug;
|
| 512 |
+
};
|
| 513 |
+
/* End XCConfigurationList section */
|
| 514 |
+
};
|
| 515 |
+
rootObject = D6A8CACD08FF01EA2A8633B6 /* Project object */;
|
| 516 |
+
}
|
FlashMoE-iOS/FlashMoEEngine/FlashMoEEngine.h
ADDED
|
@@ -0,0 +1,157 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
//
|
| 2 |
+
// FlashMoEEngine.h
|
| 3 |
+
// Flash-MoE inference engine — C API for iOS integration
|
| 4 |
+
//
|
| 5 |
+
// Wraps the Metal compute pipeline for Qwen3.5 MoE models.
|
| 6 |
+
// Expert weights stream from SSD via pread(); only K active experts
|
| 7 |
+
// are loaded per layer. Works with 4-bit and tiered (4-bit/2-bit) quantization.
|
| 8 |
+
//
|
| 9 |
+
|
| 10 |
+
#ifndef FLASHMOE_ENGINE_H
|
| 11 |
+
#define FLASHMOE_ENGINE_H
|
| 12 |
+
|
| 13 |
+
#include <stddef.h>
|
| 14 |
+
#include <stdint.h>
|
| 15 |
+
|
| 16 |
+
#ifdef __cplusplus
|
| 17 |
+
extern "C" {
|
| 18 |
+
#endif
|
| 19 |
+
|
| 20 |
+
// ---- Opaque engine handle ----
|
| 21 |
+
typedef struct FlashMoEContext FlashMoEContext;
|
| 22 |
+
|
| 23 |
+
// ---- Token callback ----
|
| 24 |
+
// Called for each generated token. Return 0 to continue, non-zero to stop.
|
| 25 |
+
typedef int (*FlashMoETokenCallback)(
|
| 26 |
+
const char *token_text, // Decoded token string (UTF-8)
|
| 27 |
+
int token_id, // Raw token ID
|
| 28 |
+
int tokens_generated, // Total tokens generated so far
|
| 29 |
+
double tokens_per_second, // Current tok/s
|
| 30 |
+
void *user_data // User context pointer
|
| 31 |
+
);
|
| 32 |
+
|
| 33 |
+
// ---- Configuration ----
|
| 34 |
+
typedef struct {
|
| 35 |
+
const char *model_path; // Path to model directory (contains config.json, packed_experts/, etc.)
|
| 36 |
+
int max_context; // Max sequence length (0 = use model default)
|
| 37 |
+
int think_budget; // Max thinking tokens (0 = unlimited)
|
| 38 |
+
int use_tiered; // 1 = use tiered quantization if available, 0 = auto-detect
|
| 39 |
+
int use_2bit; // 1 = force 2-bit experts, 0 = auto-detect from directory
|
| 40 |
+
int active_experts_k; // Override K (experts per token). 0 = use model default.
|
| 41 |
+
// Reducing K cuts I/O proportionally (e.g. K=4 on a K=10 model
|
| 42 |
+
// reduces expert reads by 60%). Quality degrades gracefully.
|
| 43 |
+
int cache_io_split; // >1 = split each expert pread into N page-aligned chunks (fanout), 0/1 = disabled
|
| 44 |
+
int cmd_merge; // 1 = merge CMD1+CMD2 for linear attention (default), 0 = separate (safer)
|
| 45 |
+
int fused_attention; // 1 = fused online softmax attention (experimental), 0 = 3-kernel fallback
|
| 46 |
+
int expert_prefetch; // 1 = prefetch next layer's experts during CMD3 (default), 0 = disabled
|
| 47 |
+
int fused_expert; // 1 = fused gate+up+SwiGLU kernel (default), 0 = separate dispatches
|
| 48 |
+
int fp16_accumulation; // 1 = use half-precision accumulation in dequant kernels (experimental), 0 = float32
|
| 49 |
+
int fp8_kv_cache; // 1 = FP8 E4M3 KV cache (4x less memory, longer context), 0 = float32
|
| 50 |
+
int sliding_window; // >0 = sliding window size for full attention layers, 0 = unlimited
|
| 51 |
+
int h2o_budget; // >0 = H2O KV cache budget (sinks + recent + heavy hitters), 0 = disabled
|
| 52 |
+
int rope_scaling_mode; // 0 = off, 1 = linear, 2 = NTK-aware, 3 = YaRN
|
| 53 |
+
float rope_scale_factor; // context extension factor (e.g. 2.0 = 2x context, 4.0 = 4x)
|
| 54 |
+
int prefill_batch; // >1 = batched GEMM prefill (e.g. 8, 16, 32). 0/1 = token-by-token (default)
|
| 55 |
+
int verbose; // 1 = log to stderr, 0 = quiet
|
| 56 |
+
} FlashMoEConfig;
|
| 57 |
+
|
| 58 |
+
// ---- Engine stats ----
|
| 59 |
+
typedef struct {
|
| 60 |
+
// Model info
|
| 61 |
+
char model_name[256];
|
| 62 |
+
int num_layers;
|
| 63 |
+
int num_experts;
|
| 64 |
+
int active_experts_k;
|
| 65 |
+
int hidden_dim;
|
| 66 |
+
int vocab_size;
|
| 67 |
+
|
| 68 |
+
// Generation stats
|
| 69 |
+
double tokens_per_second;
|
| 70 |
+
int tokens_generated;
|
| 71 |
+
double total_time_ms;
|
| 72 |
+
double ttft_ms; // Time to first token
|
| 73 |
+
|
| 74 |
+
// Memory
|
| 75 |
+
size_t weight_file_bytes; // Non-expert weights (mmap'd)
|
| 76 |
+
size_t expert_file_bytes; // Total expert data on disk
|
| 77 |
+
size_t metal_buffer_bytes; // GPU buffer allocation
|
| 78 |
+
} FlashMoEStats;
|
| 79 |
+
|
| 80 |
+
// ---- Lifecycle ----
|
| 81 |
+
|
| 82 |
+
// Create engine context. Does NOT load the model yet.
|
| 83 |
+
FlashMoEContext *flashmoe_create(void);
|
| 84 |
+
|
| 85 |
+
// Load a model from the given config. Returns 0 on success, -1 on error.
|
| 86 |
+
// This allocates Metal resources, mmaps weight files, opens expert file descriptors.
|
| 87 |
+
int flashmoe_load(FlashMoEContext *ctx, const FlashMoEConfig *config);
|
| 88 |
+
|
| 89 |
+
// Unload the current model, releasing all resources.
|
| 90 |
+
void flashmoe_unload(FlashMoEContext *ctx);
|
| 91 |
+
|
| 92 |
+
// Destroy the engine context.
|
| 93 |
+
void flashmoe_destroy(FlashMoEContext *ctx);
|
| 94 |
+
|
| 95 |
+
// ---- Generation ----
|
| 96 |
+
|
| 97 |
+
// Generate tokens from a prompt. Blocks until generation completes or is cancelled.
|
| 98 |
+
// The callback is called for each token on the calling thread.
|
| 99 |
+
// Returns the number of tokens generated, or -1 on error.
|
| 100 |
+
int flashmoe_generate(
|
| 101 |
+
FlashMoEContext *ctx,
|
| 102 |
+
const char *prompt,
|
| 103 |
+
int max_tokens,
|
| 104 |
+
FlashMoETokenCallback callback,
|
| 105 |
+
void *user_data
|
| 106 |
+
);
|
| 107 |
+
|
| 108 |
+
// Generate continuation — reuses KV cache from previous turns.
|
| 109 |
+
// Only processes the new user turn, skipping re-prefill of history.
|
| 110 |
+
// The user_content should be raw text (not formatted with chat template).
|
| 111 |
+
int flashmoe_generate_continuation(
|
| 112 |
+
FlashMoEContext *ctx,
|
| 113 |
+
const char *user_content,
|
| 114 |
+
int max_tokens,
|
| 115 |
+
FlashMoETokenCallback callback,
|
| 116 |
+
void *user_data
|
| 117 |
+
);
|
| 118 |
+
|
| 119 |
+
// Cancel an in-progress generation. Safe to call from any thread.
|
| 120 |
+
void flashmoe_cancel(FlashMoEContext *ctx);
|
| 121 |
+
|
| 122 |
+
// Reset conversation state (KV cache, linear attention state, position).
|
| 123 |
+
void flashmoe_reset(FlashMoEContext *ctx);
|
| 124 |
+
|
| 125 |
+
// ---- Stats ----
|
| 126 |
+
|
| 127 |
+
// Get current engine stats. Fills the provided struct.
|
| 128 |
+
void flashmoe_get_stats(FlashMoEContext *ctx, FlashMoEStats *stats);
|
| 129 |
+
|
| 130 |
+
// ---- Runtime Config (no model reload needed) ----
|
| 131 |
+
|
| 132 |
+
// Change inference settings without reloading the model.
|
| 133 |
+
// Sets C globals directly. K is clamped to [1, model_default].
|
| 134 |
+
void flashmoe_set_runtime_config(FlashMoEContext *ctx,
|
| 135 |
+
int active_experts_k,
|
| 136 |
+
int cmd_merge,
|
| 137 |
+
int fused_attention,
|
| 138 |
+
int cache_io_split,
|
| 139 |
+
int fp16_accumulation);
|
| 140 |
+
|
| 141 |
+
// ---- Utility ----
|
| 142 |
+
|
| 143 |
+
// Check if a model directory is valid (has config.json, packed_experts/, etc.)
|
| 144 |
+
// Returns 0 if valid, -1 if not.
|
| 145 |
+
int flashmoe_validate_model(const char *model_path);
|
| 146 |
+
|
| 147 |
+
// Get the current turn count (0 = no history, >0 = can use continuation).
|
| 148 |
+
int flashmoe_turn_count(FlashMoEContext *ctx);
|
| 149 |
+
|
| 150 |
+
// Get a human-readable error string for the last error.
|
| 151 |
+
const char *flashmoe_last_error(FlashMoEContext *ctx);
|
| 152 |
+
|
| 153 |
+
#ifdef __cplusplus
|
| 154 |
+
}
|
| 155 |
+
#endif
|
| 156 |
+
|
| 157 |
+
#endif // FLASHMOE_ENGINE_H
|
FlashMoE-iOS/FlashMoEEngine/FlashMoEEngine.m
ADDED
|
@@ -0,0 +1,1372 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* FlashMoEEngine.m — iOS wrapper for the Flash-MoE inference engine
|
| 3 |
+
*
|
| 4 |
+
* Unity build: includes infer.m directly (with CHAT_MODE to suppress main()).
|
| 5 |
+
* Provides the C API defined in FlashMoEEngine.h for Swift/SwiftUI integration.
|
| 6 |
+
*
|
| 7 |
+
* Single-instance design: iOS memory constraints mean only one model at a time.
|
| 8 |
+
* The FlashMoEContext struct holds all state, wrapping infer.m's static globals.
|
| 9 |
+
*/
|
| 10 |
+
|
| 11 |
+
#define CHAT_MODE 1 // suppress main() in infer.m
|
| 12 |
+
|
| 13 |
+
// Unity build — include the entire inference engine
|
| 14 |
+
// This gives us access to all static functions and globals
|
| 15 |
+
#include "../../metal_infer/infer.m"
|
| 16 |
+
|
| 17 |
+
#include "FlashMoEEngine.h"
|
| 18 |
+
#include <stdatomic.h>
|
| 19 |
+
#include <os/proc.h>
|
| 20 |
+
#if TARGET_OS_IPHONE
|
| 21 |
+
#import <UIKit/UIKit.h>
|
| 22 |
+
#endif
|
| 23 |
+
|
| 24 |
+
// ============================================================================
|
| 25 |
+
// FlashMoEContext — wraps engine state for the public C API
|
| 26 |
+
// ============================================================================
|
| 27 |
+
|
| 28 |
+
struct FlashMoEContext {
|
| 29 |
+
// Lifecycle state
|
| 30 |
+
int loaded; // 1 if a model is loaded
|
| 31 |
+
atomic_int cancelled; // 1 if generation should stop
|
| 32 |
+
|
| 33 |
+
// Model resources (owned)
|
| 34 |
+
WeightFile *wf;
|
| 35 |
+
Vocabulary *vocab;
|
| 36 |
+
int *layer_fds; // [num_layers] file descriptors for expert layers
|
| 37 |
+
int *layer_fds_cold_local; // [num_layers] cold file descriptors
|
| 38 |
+
void **layer_mmaps; // [num_layers] mmap'd expert data
|
| 39 |
+
size_t *layer_mmap_sizes; // [num_layers] mmap sizes
|
| 40 |
+
void **layer_states; // [num_layers] linear attention state
|
| 41 |
+
KVCache **kv_caches; // [num_layers] KV caches for full attention
|
| 42 |
+
float *hidden; // [hidden_dim] working buffer
|
| 43 |
+
float *logits; // [vocab_size] logits buffer
|
| 44 |
+
uint16_t *final_norm_w; // pointer into wf (not owned)
|
| 45 |
+
int K; // num experts per token
|
| 46 |
+
|
| 47 |
+
// Conversation state (for KV cache reuse)
|
| 48 |
+
int current_pos; // sequence position for RoPE (persists across turns)
|
| 49 |
+
int turn_count; // 0 = fresh session, >0 = has history
|
| 50 |
+
|
| 51 |
+
// Generation stats
|
| 52 |
+
double tokens_per_second;
|
| 53 |
+
int tokens_generated;
|
| 54 |
+
double total_time_ms;
|
| 55 |
+
double ttft_ms;
|
| 56 |
+
|
| 57 |
+
// Memory pressure monitoring (iOS)
|
| 58 |
+
dispatch_source_t memory_pressure_source;
|
| 59 |
+
#if TARGET_OS_IPHONE
|
| 60 |
+
id memory_warning_observer;
|
| 61 |
+
#endif
|
| 62 |
+
|
| 63 |
+
// Error state
|
| 64 |
+
char last_error[512];
|
| 65 |
+
};
|
| 66 |
+
|
| 67 |
+
// ============================================================================
|
| 68 |
+
// Shader loading for iOS — find shaders.metal in the app bundle
|
| 69 |
+
// ============================================================================
|
| 70 |
+
|
| 71 |
+
// Override the shader search path for iOS: look in the app bundle first
|
| 72 |
+
static NSString *flashmoe_find_shader_source(void) {
|
| 73 |
+
NSError *error = nil;
|
| 74 |
+
NSString *src = nil;
|
| 75 |
+
|
| 76 |
+
// 1. Try app bundle (iOS deployment)
|
| 77 |
+
NSString *bundlePath = [[NSBundle mainBundle] pathForResource:@"shaders" ofType:@"metal"];
|
| 78 |
+
if (bundlePath) {
|
| 79 |
+
src = [NSString stringWithContentsOfFile:bundlePath encoding:NSUTF8StringEncoding error:&error];
|
| 80 |
+
if (src) return src;
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
// 2. Try relative paths (macOS development / testing)
|
| 84 |
+
NSArray *paths = @[@"shaders.metal", @"metal_infer/shaders.metal"];
|
| 85 |
+
for (NSString *p in paths) {
|
| 86 |
+
src = [NSString stringWithContentsOfFile:p encoding:NSUTF8StringEncoding error:&error];
|
| 87 |
+
if (src) return src;
|
| 88 |
+
}
|
| 89 |
+
|
| 90 |
+
return nil;
|
| 91 |
+
}
|
| 92 |
+
|
| 93 |
+
// ============================================================================
|
| 94 |
+
// Public API Implementation
|
| 95 |
+
// ============================================================================
|
| 96 |
+
|
| 97 |
+
FlashMoEContext *flashmoe_create(void) {
|
| 98 |
+
FlashMoEContext *ctx = calloc(1, sizeof(FlashMoEContext));
|
| 99 |
+
if (!ctx) return NULL;
|
| 100 |
+
ctx->loaded = 0;
|
| 101 |
+
atomic_store(&ctx->cancelled, 0);
|
| 102 |
+
ctx->last_error[0] = '\0';
|
| 103 |
+
return ctx;
|
| 104 |
+
}
|
| 105 |
+
|
| 106 |
+
int flashmoe_load(FlashMoEContext *ctx, const FlashMoEConfig *config) {
|
| 107 |
+
if (!ctx || !config || !config->model_path) {
|
| 108 |
+
if (ctx) snprintf(ctx->last_error, sizeof(ctx->last_error), "Invalid arguments");
|
| 109 |
+
return -1;
|
| 110 |
+
}
|
| 111 |
+
|
| 112 |
+
// Unload any previously loaded model
|
| 113 |
+
if (ctx->loaded) {
|
| 114 |
+
flashmoe_unload(ctx);
|
| 115 |
+
}
|
| 116 |
+
|
| 117 |
+
@autoreleasepool {
|
| 118 |
+
const char *model_path = config->model_path;
|
| 119 |
+
|
| 120 |
+
// ---- Load model configuration ----
|
| 121 |
+
load_model_config(model_path);
|
| 122 |
+
alloc_tracking_arrays();
|
| 123 |
+
|
| 124 |
+
// Apply config overrides — cap context length for iOS memory constraints
|
| 125 |
+
if (config->max_context > 0) {
|
| 126 |
+
cfg.max_seq_len = config->max_context;
|
| 127 |
+
}
|
| 128 |
+
// iOS: adaptive context length based on available device memory
|
| 129 |
+
// Must account for: weight file (mmap'd), Metal buffers, delta-net state, KV caches
|
| 130 |
+
{
|
| 131 |
+
#if TARGET_OS_IPHONE || TARGET_OS_IOS
|
| 132 |
+
size_t avail = os_proc_available_memory();
|
| 133 |
+
#else
|
| 134 |
+
// macOS fallback: estimate available memory using host_statistics64
|
| 135 |
+
mach_port_t host = mach_host_self();
|
| 136 |
+
vm_size_t pageSize = 0;
|
| 137 |
+
host_page_size(host, &pageSize);
|
| 138 |
+
vm_statistics64_data_t vmStats = {0};
|
| 139 |
+
mach_msg_type_number_t count = HOST_VM_INFO64_COUNT;
|
| 140 |
+
kern_return_t kr = host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&vmStats, &count);
|
| 141 |
+
size_t avail = 0;
|
| 142 |
+
if (kr == KERN_SUCCESS) {
|
| 143 |
+
uint64_t freePages = vmStats.free_count;
|
| 144 |
+
uint64_t inactivePages = vmStats.inactive_count;
|
| 145 |
+
uint64_t speculativePages = vmStats.speculative_count;
|
| 146 |
+
avail = (size_t)((freePages + inactivePages + speculativePages) * (uint64_t)pageSize);
|
| 147 |
+
}
|
| 148 |
+
#endif
|
| 149 |
+
size_t bytes_per_elem = g_use_fp8_kv ? 1 : sizeof(float);
|
| 150 |
+
size_t kv_cost_per_pos = (size_t)cfg.num_kv_heads * cfg.head_dim * bytes_per_elem
|
| 151 |
+
* 2 // k + v
|
| 152 |
+
* cfg.num_full_attn_layers
|
| 153 |
+
* 2; // CPU + GPU mirror
|
| 154 |
+
|
| 155 |
+
// Estimate non-KV Metal memory:
|
| 156 |
+
// delta-net: num_linear_layers * v_heads * v_dim * k_dim * 4
|
| 157 |
+
// multi-expert: MAX_K * 2 * expert_size
|
| 158 |
+
// working buffers: ~50 MB
|
| 159 |
+
size_t delta_net_bytes = (size_t)cfg.num_linear_layers *
|
| 160 |
+
cfg.linear_num_v_heads * cfg.linear_value_dim * cfg.linear_key_dim * sizeof(float);
|
| 161 |
+
size_t expert_buf_bytes = (size_t)MAX_K * 2 * cfg.expert_size_4bit;
|
| 162 |
+
size_t fixed_metal = delta_net_bytes + expert_buf_bytes + 50 * 1024 * 1024;
|
| 163 |
+
|
| 164 |
+
// Reserve memory for: fixed Metal + OS headroom (2 GB) + expert page cache (at least 1 GB)
|
| 165 |
+
size_t reserved = fixed_metal + (size_t)3 * 1024 * 1024 * 1024;
|
| 166 |
+
size_t kv_budget = (avail > reserved) ? (avail - reserved) / 2 : avail / 8;
|
| 167 |
+
|
| 168 |
+
int adaptive_max = (kv_cost_per_pos > 0) ? (int)(kv_budget / kv_cost_per_pos) : 8192;
|
| 169 |
+
// Clamp to powers of 2: 512, 1024, 2048, 4096, 8192
|
| 170 |
+
int capped = 512;
|
| 171 |
+
for (int p = 512; p <= 8192; p *= 2) {
|
| 172 |
+
if (p <= adaptive_max) capped = p;
|
| 173 |
+
}
|
| 174 |
+
if (cfg.max_seq_len > capped) {
|
| 175 |
+
NSLog(@"[FlashMoE] Adaptive context: %d → %d (%.0f MB available, %.0f MB fixed Metal, KV %.0f bytes/pos)",
|
| 176 |
+
cfg.max_seq_len, capped, avail / 1e6, fixed_metal / 1e6, (double)kv_cost_per_pos);
|
| 177 |
+
cfg.max_seq_len = capped;
|
| 178 |
+
}
|
| 179 |
+
}
|
| 180 |
+
// Set runtime KV sequence limit — kv_cache_new() and GPU buffers use this
|
| 181 |
+
g_kv_seq_len = cfg.max_seq_len;
|
| 182 |
+
|
| 183 |
+
// think_budget: >0 = max thinking tokens, 0 = unlimited, -1 = disable thinking entirely
|
| 184 |
+
g_think_budget = config->think_budget;
|
| 185 |
+
|
| 186 |
+
// Set tiered mode
|
| 187 |
+
g_use_tiered = config->use_tiered;
|
| 188 |
+
g_use_2bit = config->use_2bit;
|
| 189 |
+
|
| 190 |
+
// Auto-detect 2-bit experts if not explicitly set and no 4-bit/tiered found
|
| 191 |
+
if (!g_use_2bit && !g_use_tiered) {
|
| 192 |
+
char probe_4bit[1024], probe_2bit[1024];
|
| 193 |
+
snprintf(probe_4bit, sizeof(probe_4bit), "%s/packed_experts/layer_00.bin", model_path);
|
| 194 |
+
snprintf(probe_2bit, sizeof(probe_2bit), "%s/packed_experts_2bit/layer_00.bin", model_path);
|
| 195 |
+
if (access(probe_4bit, R_OK) != 0 && access(probe_2bit, R_OK) == 0) {
|
| 196 |
+
g_use_2bit = 1;
|
| 197 |
+
NSLog(@"[FlashMoE] Auto-detected 2-bit expert files");
|
| 198 |
+
}
|
| 199 |
+
}
|
| 200 |
+
|
| 201 |
+
// Set cache I/O split (fanout mode): >1 = split expert preads into N chunks
|
| 202 |
+
if (config->cache_io_split > 1) {
|
| 203 |
+
g_cache_io_split = config->cache_io_split;
|
| 204 |
+
} else {
|
| 205 |
+
g_cache_io_split = 1; // disabled by default
|
| 206 |
+
}
|
| 207 |
+
|
| 208 |
+
// CMD1+CMD2 merge toggle (default: enabled)
|
| 209 |
+
g_cmd_merge_enabled = config->cmd_merge;
|
| 210 |
+
|
| 211 |
+
// Fused attention toggle (default: disabled, experimental)
|
| 212 |
+
g_fused_attention_enabled = config->fused_attention;
|
| 213 |
+
|
| 214 |
+
// K = experts per token from config
|
| 215 |
+
// K override: allow reducing active experts for memory-constrained devices.
|
| 216 |
+
// Lower K = less I/O per token (linear reduction). Quality degrades gracefully
|
| 217 |
+
// because the router still picks the best K experts from the full vocabulary.
|
| 218 |
+
if (config->active_experts_k > 0 && config->active_experts_k <= cfg.num_experts_per_tok) {
|
| 219 |
+
ctx->K = config->active_experts_k;
|
| 220 |
+
if (config->verbose) {
|
| 221 |
+
NSLog(@"[FlashMoE] K override: %d (model default: %d) — %.0f%% I/O reduction",
|
| 222 |
+
ctx->K, cfg.num_experts_per_tok,
|
| 223 |
+
(1.0 - (double)ctx->K / cfg.num_experts_per_tok) * 100);
|
| 224 |
+
}
|
| 225 |
+
} else {
|
| 226 |
+
ctx->K = cfg.num_experts_per_tok;
|
| 227 |
+
}
|
| 228 |
+
|
| 229 |
+
// Safety: cap K to MAX_K to prevent buffer overflow on multi-expert buffers
|
| 230 |
+
if (ctx->K > MAX_K) {
|
| 231 |
+
NSLog(@"[FlashMoE] WARNING: K=%d exceeds MAX_K=%d, capping to %d", ctx->K, MAX_K, MAX_K);
|
| 232 |
+
ctx->K = MAX_K;
|
| 233 |
+
}
|
| 234 |
+
|
| 235 |
+
// ---- Build file paths ----
|
| 236 |
+
char weights_path[1024], manifest_path[1024], vocab_path[1024];
|
| 237 |
+
|
| 238 |
+
// On iOS, weight files are in the model directory
|
| 239 |
+
snprintf(weights_path, sizeof(weights_path), "%s/model_weights.bin", model_path);
|
| 240 |
+
snprintf(manifest_path, sizeof(manifest_path), "%s/model_weights.json", model_path);
|
| 241 |
+
|
| 242 |
+
// Vocab/tokenizer: try model dir first, then app bundle
|
| 243 |
+
snprintf(vocab_path, sizeof(vocab_path), "%s/vocab.bin", model_path);
|
| 244 |
+
if (access(vocab_path, R_OK) != 0) {
|
| 245 |
+
// Try app bundle
|
| 246 |
+
NSString *bundleVocab = [[NSBundle mainBundle] pathForResource:@"vocab" ofType:@"bin"];
|
| 247 |
+
if (bundleVocab) {
|
| 248 |
+
strlcpy(vocab_path, [bundleVocab UTF8String], sizeof(vocab_path));
|
| 249 |
+
}
|
| 250 |
+
}
|
| 251 |
+
|
| 252 |
+
// ---- Set GPU-affecting flags BEFORE metal_setup() ----
|
| 253 |
+
// These flags affect buffer allocation sizes in metal_setup():
|
| 254 |
+
g_use_fp8_kv = config->fp8_kv_cache; // FP8 KV = 1 byte vs 4 bytes per element
|
| 255 |
+
g_sliding_window = config->sliding_window; // Caps KV buffer capacity
|
| 256 |
+
g_h2o_budget = config->h2o_budget; // H2O eviction budget (0 = disabled)
|
| 257 |
+
|
| 258 |
+
// RoPE scaling for context extension
|
| 259 |
+
g_rope_scaling_mode = config->rope_scaling_mode;
|
| 260 |
+
g_rope_scale_factor = config->rope_scale_factor > 0.0f ? config->rope_scale_factor : 1.0f;
|
| 261 |
+
|
| 262 |
+
// Batched prefill
|
| 263 |
+
if (config->prefill_batch > 1) {
|
| 264 |
+
g_prefill_batch = config->prefill_batch;
|
| 265 |
+
}
|
| 266 |
+
|
| 267 |
+
// ---- Initialize Metal ----
|
| 268 |
+
g_metal = metal_setup();
|
| 269 |
+
if (!g_metal) {
|
| 270 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Metal initialization failed");
|
| 271 |
+
return -1;
|
| 272 |
+
}
|
| 273 |
+
|
| 274 |
+
// Use Metal recommended working set as additional constraint on context length
|
| 275 |
+
{
|
| 276 |
+
size_t metal_budget = g_metal->recommended_working_set;
|
| 277 |
+
size_t kv_cost_per_pos_gpu = (size_t)cfg.num_kv_heads * cfg.head_dim * sizeof(float)
|
| 278 |
+
* 2 // k + v
|
| 279 |
+
* cfg.num_full_attn_layers;
|
| 280 |
+
// Estimate non-KV Metal usage (delta-net state + expert buffers + working buffers)
|
| 281 |
+
size_t delta_net_bytes = (size_t)cfg.num_linear_layers *
|
| 282 |
+
cfg.linear_num_v_heads * cfg.linear_value_dim * cfg.linear_key_dim * sizeof(float);
|
| 283 |
+
size_t expert_buf_bytes = (size_t)MAX_K * 2 * cfg.expert_size_4bit;
|
| 284 |
+
size_t fixed_gpu = delta_net_bytes + expert_buf_bytes + 50 * 1024 * 1024;
|
| 285 |
+
if (metal_budget > fixed_gpu && kv_cost_per_pos_gpu > 0) {
|
| 286 |
+
int metal_max = (int)((metal_budget - fixed_gpu) / kv_cost_per_pos_gpu);
|
| 287 |
+
// Clamp to powers of 2
|
| 288 |
+
int metal_capped = 512;
|
| 289 |
+
for (int p = 512; p <= 8192; p *= 2) {
|
| 290 |
+
if (p <= metal_max) metal_capped = p;
|
| 291 |
+
}
|
| 292 |
+
if (cfg.max_seq_len > metal_capped) {
|
| 293 |
+
NSLog(@"[FlashMoE] Metal working set constraint: %d → %d (%.1f GB budget, %.1f MB fixed GPU)",
|
| 294 |
+
cfg.max_seq_len, metal_capped, metal_budget / 1e9, fixed_gpu / 1e6);
|
| 295 |
+
cfg.max_seq_len = metal_capped;
|
| 296 |
+
g_kv_seq_len = cfg.max_seq_len;
|
| 297 |
+
}
|
| 298 |
+
}
|
| 299 |
+
}
|
| 300 |
+
|
| 301 |
+
// Set up memory pressure monitoring
|
| 302 |
+
{
|
| 303 |
+
dispatch_source_t src = dispatch_source_create(
|
| 304 |
+
DISPATCH_SOURCE_TYPE_MEMORYPRESSURE, 0,
|
| 305 |
+
DISPATCH_MEMORYPRESSURE_WARN | DISPATCH_MEMORYPRESSURE_CRITICAL,
|
| 306 |
+
dispatch_get_main_queue());
|
| 307 |
+
if (src) {
|
| 308 |
+
FlashMoEContext *ctx_capture = ctx;
|
| 309 |
+
dispatch_source_set_event_handler(src, ^{
|
| 310 |
+
unsigned long status = dispatch_source_get_data(src);
|
| 311 |
+
if (status & DISPATCH_MEMORYPRESSURE_CRITICAL) {
|
| 312 |
+
NSLog(@"[FlashMoE] CRITICAL memory pressure — cancelling generation");
|
| 313 |
+
if (ctx_capture) {
|
| 314 |
+
atomic_store(&ctx_capture->cancelled, 1);
|
| 315 |
+
}
|
| 316 |
+
} else if (status & DISPATCH_MEMORYPRESSURE_WARN) {
|
| 317 |
+
NSLog(@"[FlashMoE] WARNING: memory pressure elevated");
|
| 318 |
+
}
|
| 319 |
+
});
|
| 320 |
+
dispatch_resume(src);
|
| 321 |
+
ctx->memory_pressure_source = src;
|
| 322 |
+
}
|
| 323 |
+
}
|
| 324 |
+
|
| 325 |
+
// ---- Initialize I/O thread pool ----
|
| 326 |
+
io_pool_init();
|
| 327 |
+
|
| 328 |
+
// ---- Load weights ----
|
| 329 |
+
ctx->wf = open_weights(weights_path, manifest_path);
|
| 330 |
+
if (!ctx->wf) {
|
| 331 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to load weights from %s", weights_path);
|
| 332 |
+
return -1;
|
| 333 |
+
}
|
| 334 |
+
|
| 335 |
+
// Wrap weight file for Metal GPU access
|
| 336 |
+
if (ctx->wf->is_split) {
|
| 337 |
+
metal_set_weights_split(g_metal, ctx->wf);
|
| 338 |
+
} else {
|
| 339 |
+
metal_set_weights(g_metal, ctx->wf->data, ctx->wf->size);
|
| 340 |
+
}
|
| 341 |
+
|
| 342 |
+
// ---- Load vocabulary ----
|
| 343 |
+
ctx->vocab = load_vocab(vocab_path);
|
| 344 |
+
if (!ctx->vocab) {
|
| 345 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to load vocabulary from %s", vocab_path);
|
| 346 |
+
return -1;
|
| 347 |
+
}
|
| 348 |
+
|
| 349 |
+
// ---- Initialize tokenizer ----
|
| 350 |
+
init_tokenizer();
|
| 351 |
+
|
| 352 |
+
// ---- Auto-detect/load tiered manifest ----
|
| 353 |
+
if (!g_use_2bit && !g_use_tiered) {
|
| 354 |
+
char probe[1024];
|
| 355 |
+
snprintf(probe, sizeof(probe), "%s/packed_experts_tiered/tiered_manifest.json", model_path);
|
| 356 |
+
if (access(probe, F_OK) == 0) {
|
| 357 |
+
if (load_tiered_manifest(model_path)) {
|
| 358 |
+
g_use_tiered = 1;
|
| 359 |
+
}
|
| 360 |
+
}
|
| 361 |
+
}
|
| 362 |
+
if (g_use_tiered && !g_tiered_manifest) {
|
| 363 |
+
if (!load_tiered_manifest(model_path)) {
|
| 364 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error),
|
| 365 |
+
"Tiered mode requested but no manifest found");
|
| 366 |
+
return -1;
|
| 367 |
+
}
|
| 368 |
+
}
|
| 369 |
+
|
| 370 |
+
// ---- Open packed expert files ----
|
| 371 |
+
ctx->layer_fds = calloc(cfg.num_layers, sizeof(int));
|
| 372 |
+
ctx->layer_fds_cold_local = calloc(cfg.num_layers, sizeof(int));
|
| 373 |
+
ctx->layer_mmaps = calloc(cfg.num_layers, sizeof(void *));
|
| 374 |
+
ctx->layer_mmap_sizes = calloc(cfg.num_layers, sizeof(size_t));
|
| 375 |
+
if (!ctx->layer_fds || !ctx->layer_fds_cold_local || !ctx->layer_mmaps || !ctx->layer_mmap_sizes) {
|
| 376 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to allocate layer arrays (%d layers)", cfg.num_layers);
|
| 377 |
+
return -1;
|
| 378 |
+
}
|
| 379 |
+
|
| 380 |
+
memset(g_expert_seen, 0, cfg.num_layers * ((cfg.num_experts + 7) / 8));
|
| 381 |
+
|
| 382 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 383 |
+
char path[1024];
|
| 384 |
+
snprintf(path, sizeof(path), "%s/%s/layer_%02d.bin", model_path,
|
| 385 |
+
g_use_tiered ? "packed_experts_tiered" :
|
| 386 |
+
g_use_2bit ? "packed_experts_2bit" : "packed_experts", i);
|
| 387 |
+
ctx->layer_fds[i] = open(path, O_RDONLY);
|
| 388 |
+
ctx->layer_fds_cold_local[i] = -1;
|
| 389 |
+
ctx->layer_mmaps[i] = MAP_FAILED;
|
| 390 |
+
ctx->layer_mmap_sizes[i] = 0;
|
| 391 |
+
if (ctx->layer_fds[i] >= 0) {
|
| 392 |
+
fcntl(ctx->layer_fds[i], F_RDAHEAD, 0);
|
| 393 |
+
#if TARGET_OS_IOS
|
| 394 |
+
// On real iOS devices, do NOT mmap expert files.
|
| 395 |
+
// mmap'ing all expert layers (e.g. 60 × 1.9GB = 112GB for 397B)
|
| 396 |
+
// causes jetsam kills. Use pread() only on iOS.
|
| 397 |
+
// (macOS / Mac Catalyst can still mmap.)
|
| 398 |
+
if (![[NSProcessInfo processInfo] isMacCatalystApp]) {
|
| 399 |
+
// pread-only: leave layer_mmaps[i] = MAP_FAILED
|
| 400 |
+
} else
|
| 401 |
+
#endif
|
| 402 |
+
{
|
| 403 |
+
struct stat st;
|
| 404 |
+
if (fstat(ctx->layer_fds[i], &st) == 0 && st.st_size > 0) {
|
| 405 |
+
ctx->layer_mmaps[i] = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE,
|
| 406 |
+
ctx->layer_fds[i], 0);
|
| 407 |
+
if (ctx->layer_mmaps[i] != MAP_FAILED) {
|
| 408 |
+
ctx->layer_mmap_sizes[i] = st.st_size;
|
| 409 |
+
}
|
| 410 |
+
}
|
| 411 |
+
}
|
| 412 |
+
}
|
| 413 |
+
}
|
| 414 |
+
|
| 415 |
+
// Wire up global cold fds
|
| 416 |
+
g_layer_fds_cold = ctx->layer_fds_cold_local;
|
| 417 |
+
|
| 418 |
+
// Wire up cross-layer prefetch globals
|
| 419 |
+
g_layer_fds_global = ctx->layer_fds;
|
| 420 |
+
g_layer_mmaps_global = (void **)ctx->layer_mmaps;
|
| 421 |
+
g_layer_mmap_sizes_global = ctx->layer_mmap_sizes;
|
| 422 |
+
g_expert_prefetch_enabled = config->expert_prefetch;
|
| 423 |
+
g_fused_expert_enabled = config->fused_expert;
|
| 424 |
+
g_use_fp16_accum = config->fp16_accumulation;
|
| 425 |
+
g_use_fp8_kv = config->fp8_kv_cache;
|
| 426 |
+
g_sliding_window = config->sliding_window;
|
| 427 |
+
g_h2o_budget = config->h2o_budget;
|
| 428 |
+
g_rope_scaling_mode = config->rope_scaling_mode;
|
| 429 |
+
g_rope_scale_factor = config->rope_scale_factor > 0.0f ? config->rope_scale_factor : 1.0f;
|
| 430 |
+
if (config->prefill_batch > 1) g_prefill_batch = config->prefill_batch;
|
| 431 |
+
g_prefetch_active = 0;
|
| 432 |
+
g_prefetch_layer = -1;
|
| 433 |
+
g_prefetch_hits_total = 0;
|
| 434 |
+
g_prefetch_misses_total = 0;
|
| 435 |
+
|
| 436 |
+
// ---- Allocate deferred expert state ----
|
| 437 |
+
g_deferred.h_mid = calloc(cfg.hidden_dim, sizeof(float));
|
| 438 |
+
if (!g_deferred.h_mid) {
|
| 439 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to allocate deferred expert state (%d floats)", cfg.hidden_dim);
|
| 440 |
+
return -1;
|
| 441 |
+
}
|
| 442 |
+
|
| 443 |
+
// ---- Allocate per-layer state ----
|
| 444 |
+
ctx->layer_states = calloc(cfg.num_layers, sizeof(void *));
|
| 445 |
+
ctx->kv_caches = calloc(cfg.num_layers, sizeof(KVCache *));
|
| 446 |
+
if (!ctx->layer_states || !ctx->kv_caches) {
|
| 447 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to allocate per-layer state arrays (%d layers)", cfg.num_layers);
|
| 448 |
+
return -1;
|
| 449 |
+
}
|
| 450 |
+
|
| 451 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 452 |
+
if (cfg.is_full_attn[i]) {
|
| 453 |
+
ctx->kv_caches[i] = kv_cache_new();
|
| 454 |
+
} else {
|
| 455 |
+
ctx->layer_states[i] = linear_attn_state_new();
|
| 456 |
+
}
|
| 457 |
+
}
|
| 458 |
+
|
| 459 |
+
// ---- Allocate working buffers ----
|
| 460 |
+
ctx->hidden = calloc(cfg.hidden_dim, sizeof(float));
|
| 461 |
+
ctx->logits = calloc(cfg.vocab_size, sizeof(float));
|
| 462 |
+
if (!ctx->hidden || !ctx->logits) {
|
| 463 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to allocate working buffers (hidden=%d, vocab=%d)", cfg.hidden_dim, cfg.vocab_size);
|
| 464 |
+
return -1;
|
| 465 |
+
}
|
| 466 |
+
ctx->final_norm_w = get_tensor_ptr(ctx->wf, "model.norm.weight");
|
| 467 |
+
|
| 468 |
+
// ---- Build layer cache (precomputes weight pointers) ----
|
| 469 |
+
build_layer_cache(ctx->wf);
|
| 470 |
+
|
| 471 |
+
ctx->loaded = 1;
|
| 472 |
+
|
| 473 |
+
// Register for iOS memory warning notifications
|
| 474 |
+
#if TARGET_OS_IPHONE
|
| 475 |
+
{
|
| 476 |
+
FlashMoEContext *ctx_warn = ctx;
|
| 477 |
+
ctx->memory_warning_observer = [[NSNotificationCenter defaultCenter]
|
| 478 |
+
addObserverForName:UIApplicationDidReceiveMemoryWarningNotification
|
| 479 |
+
object:nil
|
| 480 |
+
queue:[NSOperationQueue mainQueue]
|
| 481 |
+
usingBlock:^(NSNotification *note) {
|
| 482 |
+
NSLog(@"[FlashMoE] iOS memory warning received — cancelling generation");
|
| 483 |
+
if (ctx_warn && ctx_warn->loaded) {
|
| 484 |
+
atomic_store(&ctx_warn->cancelled, 1);
|
| 485 |
+
}
|
| 486 |
+
}];
|
| 487 |
+
}
|
| 488 |
+
#endif
|
| 489 |
+
|
| 490 |
+
if (config->verbose) {
|
| 491 |
+
NSLog(@"[FlashMoE] Model loaded: %d layers, %d experts (K=%d), hidden=%d",
|
| 492 |
+
cfg.num_layers, cfg.num_experts, ctx->K, cfg.hidden_dim);
|
| 493 |
+
}
|
| 494 |
+
|
| 495 |
+
return 0;
|
| 496 |
+
}
|
| 497 |
+
}
|
| 498 |
+
|
| 499 |
+
void flashmoe_unload(FlashMoEContext *ctx) {
|
| 500 |
+
if (!ctx || !ctx->loaded) return;
|
| 501 |
+
|
| 502 |
+
@autoreleasepool {
|
| 503 |
+
// Wait for any in-flight GPU work
|
| 504 |
+
if (g_deferred.active) {
|
| 505 |
+
[g_deferred.cmd_experts waitUntilCompleted];
|
| 506 |
+
g_deferred.active = 0;
|
| 507 |
+
g_deferred.cmd_experts = nil;
|
| 508 |
+
}
|
| 509 |
+
|
| 510 |
+
// Wait for any in-flight async pread
|
| 511 |
+
if (g_async_pread.active) {
|
| 512 |
+
dispatch_group_wait(g_async_pread.group, DISPATCH_TIME_FOREVER);
|
| 513 |
+
g_async_pread.active = 0;
|
| 514 |
+
}
|
| 515 |
+
|
| 516 |
+
// Clear cross-layer prefetch globals
|
| 517 |
+
g_layer_fds_global = NULL;
|
| 518 |
+
g_layer_mmaps_global = NULL;
|
| 519 |
+
g_layer_mmap_sizes_global = NULL;
|
| 520 |
+
g_prefetch_active = 0;
|
| 521 |
+
g_prefetch_layer = -1;
|
| 522 |
+
|
| 523 |
+
// Shutdown I/O pool
|
| 524 |
+
io_pool_shutdown();
|
| 525 |
+
|
| 526 |
+
// Close expert files
|
| 527 |
+
if (ctx->layer_fds) {
|
| 528 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 529 |
+
if (ctx->layer_mmaps && ctx->layer_mmaps[i] != MAP_FAILED)
|
| 530 |
+
munmap(ctx->layer_mmaps[i], ctx->layer_mmap_sizes[i]);
|
| 531 |
+
if (ctx->layer_fds[i] >= 0)
|
| 532 |
+
close(ctx->layer_fds[i]);
|
| 533 |
+
if (ctx->layer_fds_cold_local && ctx->layer_fds_cold_local[i] >= 0)
|
| 534 |
+
close(ctx->layer_fds_cold_local[i]);
|
| 535 |
+
}
|
| 536 |
+
free(ctx->layer_fds); ctx->layer_fds = NULL;
|
| 537 |
+
free(ctx->layer_fds_cold_local); ctx->layer_fds_cold_local = NULL;
|
| 538 |
+
free(ctx->layer_mmaps); ctx->layer_mmaps = NULL;
|
| 539 |
+
free(ctx->layer_mmap_sizes); ctx->layer_mmap_sizes = NULL;
|
| 540 |
+
}
|
| 541 |
+
|
| 542 |
+
// Free per-layer state
|
| 543 |
+
if (ctx->layer_states) {
|
| 544 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 545 |
+
if (ctx->kv_caches && ctx->kv_caches[i])
|
| 546 |
+
kv_cache_free(ctx->kv_caches[i]);
|
| 547 |
+
if (ctx->layer_states[i])
|
| 548 |
+
linear_attn_state_free(ctx->layer_states[i]);
|
| 549 |
+
}
|
| 550 |
+
free(ctx->layer_states); ctx->layer_states = NULL;
|
| 551 |
+
free(ctx->kv_caches); ctx->kv_caches = NULL;
|
| 552 |
+
}
|
| 553 |
+
|
| 554 |
+
// Free working buffers
|
| 555 |
+
free(ctx->hidden); ctx->hidden = NULL;
|
| 556 |
+
free(ctx->logits); ctx->logits = NULL;
|
| 557 |
+
|
| 558 |
+
// Free deferred state
|
| 559 |
+
free(g_deferred.h_mid); g_deferred.h_mid = NULL;
|
| 560 |
+
|
| 561 |
+
// Free weight file (munmap + manifest)
|
| 562 |
+
if (ctx->wf) {
|
| 563 |
+
if (ctx->wf->data) munmap(ctx->wf->data, ctx->wf->size);
|
| 564 |
+
if (ctx->wf->manifest) {
|
| 565 |
+
free(ctx->wf->manifest->tensors);
|
| 566 |
+
free(ctx->wf->manifest);
|
| 567 |
+
}
|
| 568 |
+
free(ctx->wf);
|
| 569 |
+
ctx->wf = NULL;
|
| 570 |
+
}
|
| 571 |
+
ctx->final_norm_w = NULL;
|
| 572 |
+
|
| 573 |
+
// Reset tensor hash table (points into freed manifest)
|
| 574 |
+
memset(tensor_ht, 0, sizeof(tensor_ht));
|
| 575 |
+
tensor_ht_built = 0;
|
| 576 |
+
|
| 577 |
+
// Free vocabulary
|
| 578 |
+
if (ctx->vocab) {
|
| 579 |
+
free(ctx->vocab);
|
| 580 |
+
ctx->vocab = NULL;
|
| 581 |
+
}
|
| 582 |
+
|
| 583 |
+
// Free config dynamic arrays
|
| 584 |
+
free(cfg.is_full_attn); cfg.is_full_attn = NULL;
|
| 585 |
+
|
| 586 |
+
// Free tracking arrays (allocated by alloc_tracking_arrays)
|
| 587 |
+
free(g_expert_freq); g_expert_freq = NULL;
|
| 588 |
+
free(g_expert_seen); g_expert_seen = NULL;
|
| 589 |
+
free(g_lz4_index); g_lz4_index = NULL;
|
| 590 |
+
free(g_cache_seen); g_cache_seen = NULL;
|
| 591 |
+
free(g_cache_last_touch_token); g_cache_last_touch_token = NULL;
|
| 592 |
+
free(g_cache_last_evict_token); g_cache_last_evict_token = NULL;
|
| 593 |
+
free(g_pred_experts); g_pred_experts = NULL;
|
| 594 |
+
free(g_pred_count); g_pred_count = NULL;
|
| 595 |
+
|
| 596 |
+
// Reset layer cache so it rebuilds on next load
|
| 597 |
+
free(layer_cache); layer_cache = NULL;
|
| 598 |
+
layer_cache_built = 0;
|
| 599 |
+
|
| 600 |
+
// Free tiered manifest
|
| 601 |
+
if (g_tiered_manifest) {
|
| 602 |
+
free(g_tiered_manifest);
|
| 603 |
+
g_tiered_manifest = NULL;
|
| 604 |
+
g_use_tiered = 0;
|
| 605 |
+
}
|
| 606 |
+
|
| 607 |
+
// Reset prediction state
|
| 608 |
+
g_pred_enabled = 0;
|
| 609 |
+
g_pred_generating = 0;
|
| 610 |
+
g_pred_valid = 0;
|
| 611 |
+
g_pred_hits = 0;
|
| 612 |
+
g_pred_misses = 0;
|
| 613 |
+
g_pred_layers = 0;
|
| 614 |
+
|
| 615 |
+
// Reset global flags for clean reload
|
| 616 |
+
g_freq_tracking = 0;
|
| 617 |
+
g_cache_telemetry_enabled = 0;
|
| 618 |
+
g_kv_seq_len = 0;
|
| 619 |
+
|
| 620 |
+
// Remove iOS memory warning observer
|
| 621 |
+
#if TARGET_OS_IPHONE
|
| 622 |
+
if (ctx->memory_warning_observer) {
|
| 623 |
+
[[NSNotificationCenter defaultCenter] removeObserver:ctx->memory_warning_observer];
|
| 624 |
+
ctx->memory_warning_observer = nil;
|
| 625 |
+
}
|
| 626 |
+
#endif
|
| 627 |
+
|
| 628 |
+
// Cancel memory pressure monitoring
|
| 629 |
+
if (ctx->memory_pressure_source) {
|
| 630 |
+
dispatch_source_cancel(ctx->memory_pressure_source);
|
| 631 |
+
ctx->memory_pressure_source = NULL;
|
| 632 |
+
}
|
| 633 |
+
|
| 634 |
+
// Release Metal context
|
| 635 |
+
// MetalCtx is malloc'd but contains ARC-managed id<> objects.
|
| 636 |
+
// Must nil every id<> field so ARC decrements refcounts before free().
|
| 637 |
+
// Without this, switching models corrupts the heap (objc Method cache corrupted).
|
| 638 |
+
if (g_metal) {
|
| 639 |
+
// Nil every ARC-managed id<> field so refcounts are decremented.
|
| 640 |
+
// Without this, switching models leaks Metal objects and corrupts the heap.
|
| 641 |
+
g_metal->device = nil;
|
| 642 |
+
g_metal->queue = nil;
|
| 643 |
+
g_metal->library = nil;
|
| 644 |
+
// Pipeline states
|
| 645 |
+
g_metal->matvec_v3 = nil;
|
| 646 |
+
g_metal->matvec_v5 = nil;
|
| 647 |
+
g_metal->matvec_fast = nil;
|
| 648 |
+
g_metal->matvec_2bit = nil;
|
| 649 |
+
g_metal->rms_norm_sum = nil;
|
| 650 |
+
g_metal->rms_norm_apply = nil;
|
| 651 |
+
g_metal->rms_norm_apply_bf16 = nil;
|
| 652 |
+
g_metal->residual_add = nil;
|
| 653 |
+
g_metal->swiglu = nil;
|
| 654 |
+
g_metal->fused_gate_up = nil;
|
| 655 |
+
g_metal->matvec_v3_fp16 = nil;
|
| 656 |
+
g_metal->matvec_2bit_fp16 = nil;
|
| 657 |
+
g_metal->fused_gate_up_fp16 = nil;
|
| 658 |
+
g_metal->attn_scores_pipe = nil;
|
| 659 |
+
g_metal->attn_softmax_pipe = nil;
|
| 660 |
+
g_metal->attn_values_pipe = nil;
|
| 661 |
+
g_metal->sigmoid_gate_pipe = nil;
|
| 662 |
+
g_metal->moe_combine_residual = nil;
|
| 663 |
+
// GPU linear attention pipelines
|
| 664 |
+
g_metal->delta_net_step = nil;
|
| 665 |
+
g_metal->conv1d_step = nil;
|
| 666 |
+
g_metal->rms_norm_qk = nil;
|
| 667 |
+
g_metal->compute_decay_beta = nil;
|
| 668 |
+
g_metal->gated_rms_norm = nil;
|
| 669 |
+
// Shared event
|
| 670 |
+
g_metal->pipeline_event = nil;
|
| 671 |
+
// Buffers
|
| 672 |
+
g_metal->buf_input = nil;
|
| 673 |
+
g_metal->buf_output = nil;
|
| 674 |
+
g_metal->wf_buf = nil;
|
| 675 |
+
g_metal->wf_staging = nil;
|
| 676 |
+
for (int i = 0; i < MAX_WF_CHUNKS; i++) g_metal->wf_chunks[i] = nil;
|
| 677 |
+
for (int i = 0; i < MAX_BATCH_SLOTS; i++) g_metal->batch_out[i] = nil;
|
| 678 |
+
// Expert buffers
|
| 679 |
+
g_metal->buf_expert_data = nil;
|
| 680 |
+
g_metal->buf_expert_input = nil;
|
| 681 |
+
g_metal->buf_expert_gate = nil;
|
| 682 |
+
g_metal->buf_expert_up = nil;
|
| 683 |
+
g_metal->buf_expert_act = nil;
|
| 684 |
+
g_metal->buf_expert_out = nil;
|
| 685 |
+
for (int i = 0; i < MAX_K; i++) {
|
| 686 |
+
g_metal->buf_multi_expert_data[i] = nil;
|
| 687 |
+
g_metal->buf_multi_expert_data_B[i] = nil;
|
| 688 |
+
g_metal->buf_multi_expert_gate[i] = nil;
|
| 689 |
+
g_metal->buf_multi_expert_up[i] = nil;
|
| 690 |
+
g_metal->buf_multi_expert_act[i] = nil;
|
| 691 |
+
g_metal->buf_multi_expert_out[i] = nil;
|
| 692 |
+
}
|
| 693 |
+
g_metal->buf_multi_expert_input = nil;
|
| 694 |
+
g_metal->buf_shared_gate = nil;
|
| 695 |
+
g_metal->buf_shared_up = nil;
|
| 696 |
+
g_metal->buf_shared_act = nil;
|
| 697 |
+
g_metal->buf_shared_out = nil;
|
| 698 |
+
g_metal->buf_residual = nil;
|
| 699 |
+
g_metal->buf_h_mid = nil;
|
| 700 |
+
g_metal->buf_sum_sq = nil;
|
| 701 |
+
g_metal->buf_moe_hidden = nil;
|
| 702 |
+
g_metal->buf_combine_params = nil;
|
| 703 |
+
g_metal->buf_cmd3_sum_sq = nil;
|
| 704 |
+
// GPU attention buffers
|
| 705 |
+
g_metal->buf_attn_q = nil;
|
| 706 |
+
g_metal->buf_attn_scores = nil;
|
| 707 |
+
g_metal->buf_attn_out = nil;
|
| 708 |
+
g_metal->buf_attn_gate = nil;
|
| 709 |
+
if (g_metal->buf_kv_k) {
|
| 710 |
+
for (int i = 0; i < cfg.num_full_attn_layers; i++) {
|
| 711 |
+
g_metal->buf_kv_k[i] = nil;
|
| 712 |
+
g_metal->buf_kv_v[i] = nil;
|
| 713 |
+
}
|
| 714 |
+
free(g_metal->buf_kv_k); g_metal->buf_kv_k = NULL;
|
| 715 |
+
free(g_metal->buf_kv_v); g_metal->buf_kv_v = NULL;
|
| 716 |
+
}
|
| 717 |
+
// Delta-net GPU buffers
|
| 718 |
+
if (g_metal->buf_delta_state) {
|
| 719 |
+
for (int i = 0; i < cfg.num_linear_layers; i++) {
|
| 720 |
+
g_metal->buf_delta_state[i] = nil;
|
| 721 |
+
}
|
| 722 |
+
free(g_metal->buf_delta_state); g_metal->buf_delta_state = NULL;
|
| 723 |
+
}
|
| 724 |
+
if (g_metal->buf_conv_state) {
|
| 725 |
+
for (int i = 0; i < cfg.num_linear_layers; i++) {
|
| 726 |
+
g_metal->buf_conv_state[i] = nil;
|
| 727 |
+
}
|
| 728 |
+
free(g_metal->buf_conv_state); g_metal->buf_conv_state = NULL;
|
| 729 |
+
}
|
| 730 |
+
// Delta-net scratch buffers
|
| 731 |
+
g_metal->buf_delta_q = nil;
|
| 732 |
+
g_metal->buf_delta_k = nil;
|
| 733 |
+
g_metal->buf_delta_v = nil;
|
| 734 |
+
g_metal->buf_delta_g_decay = nil;
|
| 735 |
+
g_metal->buf_delta_beta = nil;
|
| 736 |
+
g_metal->buf_delta_output = nil;
|
| 737 |
+
g_metal->buf_conv_input = nil;
|
| 738 |
+
g_metal->buf_conv_output = nil;
|
| 739 |
+
free(g_metal);
|
| 740 |
+
g_metal = NULL;
|
| 741 |
+
}
|
| 742 |
+
|
| 743 |
+
ctx->loaded = 0;
|
| 744 |
+
}
|
| 745 |
+
}
|
| 746 |
+
|
| 747 |
+
void flashmoe_destroy(FlashMoEContext *ctx) {
|
| 748 |
+
if (!ctx) return;
|
| 749 |
+
flashmoe_unload(ctx);
|
| 750 |
+
free(ctx);
|
| 751 |
+
}
|
| 752 |
+
|
| 753 |
+
// ============================================================================
|
| 754 |
+
// Generation — the core inference loop adapted for callback-based streaming
|
| 755 |
+
// ============================================================================
|
| 756 |
+
|
| 757 |
+
int flashmoe_generate(
|
| 758 |
+
FlashMoEContext *ctx,
|
| 759 |
+
const char *prompt,
|
| 760 |
+
int max_tokens,
|
| 761 |
+
FlashMoETokenCallback callback,
|
| 762 |
+
void *user_data
|
| 763 |
+
) {
|
| 764 |
+
if (!ctx || !ctx->loaded || !prompt) {
|
| 765 |
+
if (ctx) snprintf(ctx->last_error, sizeof(ctx->last_error), "Engine not loaded or invalid arguments");
|
| 766 |
+
return -1;
|
| 767 |
+
}
|
| 768 |
+
|
| 769 |
+
@autoreleasepool {
|
| 770 |
+
atomic_store(&ctx->cancelled, 0);
|
| 771 |
+
ctx->tokens_generated = 0;
|
| 772 |
+
ctx->tokens_per_second = 0;
|
| 773 |
+
|
| 774 |
+
// Pre-flight memory check (iOS)
|
| 775 |
+
#if TARGET_OS_IPHONE
|
| 776 |
+
{
|
| 777 |
+
size_t avail = os_proc_available_memory();
|
| 778 |
+
if (avail < 500 * 1024 * 1024) {
|
| 779 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error),
|
| 780 |
+
"Insufficient memory (%.0f MB available, need 500+ MB)", avail / 1e6);
|
| 781 |
+
return -1;
|
| 782 |
+
}
|
| 783 |
+
}
|
| 784 |
+
#endif
|
| 785 |
+
|
| 786 |
+
double t0 = now_ms();
|
| 787 |
+
|
| 788 |
+
// ---- Tokenize prompt ----
|
| 789 |
+
PromptTokens *pt = encode_prompt_text_to_tokens(prompt);
|
| 790 |
+
if (!pt) {
|
| 791 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to tokenize prompt");
|
| 792 |
+
return -1;
|
| 793 |
+
}
|
| 794 |
+
|
| 795 |
+
int K = ctx->K;
|
| 796 |
+
|
| 797 |
+
// ---- Reset state for new generation ----
|
| 798 |
+
reset_delta_net_state();
|
| 799 |
+
// Reset KV cache lengths
|
| 800 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 801 |
+
if (ctx->kv_caches[i]) {
|
| 802 |
+
ctx->kv_caches[i]->len = 0;
|
| 803 |
+
}
|
| 804 |
+
}
|
| 805 |
+
|
| 806 |
+
int pos = 0;
|
| 807 |
+
|
| 808 |
+
// ---- Batch prefill: embed all prompt tokens ----
|
| 809 |
+
float *embed_batch = NULL;
|
| 810 |
+
if (pt->count > 1) {
|
| 811 |
+
embed_batch = malloc((size_t)pt->count * cfg.hidden_dim * sizeof(float));
|
| 812 |
+
for (int i = 0; i < pt->count; i++) {
|
| 813 |
+
embed_lookup(ctx->wf, pt->ids[i], embed_batch + (size_t)i * cfg.hidden_dim);
|
| 814 |
+
}
|
| 815 |
+
}
|
| 816 |
+
|
| 817 |
+
// ---- Prefill intermediate tokens ----
|
| 818 |
+
if (pt->count > 1) {
|
| 819 |
+
int prefill_count = pt->count - 1; // last token processed separately for full hidden state
|
| 820 |
+
int batched_done = 0;
|
| 821 |
+
|
| 822 |
+
// Try batched GEMM prefill (reads each weight row once for all tokens)
|
| 823 |
+
if (g_prefill_batch > 1 && embed_batch) {
|
| 824 |
+
batched_done = batched_prefill(ctx->wf, embed_batch, prefill_count, pos,
|
| 825 |
+
K, ctx->layer_fds, ctx->layer_mmaps,
|
| 826 |
+
ctx->kv_caches, ctx->layer_states);
|
| 827 |
+
if (batched_done > 0) {
|
| 828 |
+
pos += batched_done;
|
| 829 |
+
NSLog(@"[FlashMoE] Batched prefill: %d tokens processed", batched_done);
|
| 830 |
+
}
|
| 831 |
+
}
|
| 832 |
+
|
| 833 |
+
// Fall back to per-token for remaining tokens
|
| 834 |
+
for (int token_idx = batched_done; token_idx < prefill_count; token_idx++) { @autoreleasepool {
|
| 835 |
+
if (atomic_load(&ctx->cancelled)) {
|
| 836 |
+
free(embed_batch);
|
| 837 |
+
free(pt->ids); free(pt);
|
| 838 |
+
return ctx->tokens_generated;
|
| 839 |
+
}
|
| 840 |
+
|
| 841 |
+
memcpy(ctx->hidden, embed_batch + (size_t)token_idx * cfg.hidden_dim,
|
| 842 |
+
cfg.hidden_dim * sizeof(float));
|
| 843 |
+
|
| 844 |
+
for (int layer = 0; layer < cfg.num_layers; layer++) {
|
| 845 |
+
int is_full = cfg.is_full_attn[layer];
|
| 846 |
+
fused_layer_forward(ctx->wf, layer, ctx->hidden,
|
| 847 |
+
is_full ? ctx->kv_caches[layer] : NULL,
|
| 848 |
+
is_full ? NULL : ctx->layer_states[layer],
|
| 849 |
+
pos,
|
| 850 |
+
ctx->layer_mmaps[layer] != MAP_FAILED ? ctx->layer_mmaps[layer] : NULL,
|
| 851 |
+
K, ctx->layer_fds[layer]);
|
| 852 |
+
}
|
| 853 |
+
discard_deferred_experts();
|
| 854 |
+
pos++;
|
| 855 |
+
} /* @autoreleasepool */ }
|
| 856 |
+
}
|
| 857 |
+
|
| 858 |
+
// ---- Last prefill token (need full hidden state) ----
|
| 859 |
+
{
|
| 860 |
+
if (embed_batch) {
|
| 861 |
+
memcpy(ctx->hidden, embed_batch + (size_t)(pt->count - 1) * cfg.hidden_dim,
|
| 862 |
+
cfg.hidden_dim * sizeof(float));
|
| 863 |
+
} else {
|
| 864 |
+
embed_lookup(ctx->wf, pt->ids[0], ctx->hidden);
|
| 865 |
+
}
|
| 866 |
+
|
| 867 |
+
for (int layer = 0; layer < cfg.num_layers; layer++) {
|
| 868 |
+
int is_full = cfg.is_full_attn[layer];
|
| 869 |
+
fused_layer_forward(ctx->wf, layer, ctx->hidden,
|
| 870 |
+
is_full ? ctx->kv_caches[layer] : NULL,
|
| 871 |
+
is_full ? NULL : ctx->layer_states[layer],
|
| 872 |
+
pos,
|
| 873 |
+
ctx->layer_mmaps[layer] != MAP_FAILED ? ctx->layer_mmaps[layer] : NULL,
|
| 874 |
+
K, ctx->layer_fds[layer]);
|
| 875 |
+
}
|
| 876 |
+
complete_deferred_experts();
|
| 877 |
+
pos++;
|
| 878 |
+
}
|
| 879 |
+
|
| 880 |
+
if (embed_batch) { free(embed_batch); embed_batch = NULL; }
|
| 881 |
+
|
| 882 |
+
// ---- Final norm + LM head + sample first token ----
|
| 883 |
+
if (ctx->final_norm_w) {
|
| 884 |
+
float *normed = malloc(cfg.hidden_dim * sizeof(float));
|
| 885 |
+
cpu_rms_norm(ctx->hidden, ctx->final_norm_w, normed, cfg.hidden_dim, cfg.rms_norm_eps);
|
| 886 |
+
memcpy(ctx->hidden, normed, cfg.hidden_dim * sizeof(float));
|
| 887 |
+
free(normed);
|
| 888 |
+
}
|
| 889 |
+
|
| 890 |
+
lm_head_forward(ctx->wf, ctx->hidden, ctx->logits);
|
| 891 |
+
int next_token = cpu_argmax(ctx->logits, cfg.vocab_size);
|
| 892 |
+
|
| 893 |
+
ctx->ttft_ms = now_ms() - t0;
|
| 894 |
+
ctx->tokens_generated = 1;
|
| 895 |
+
|
| 896 |
+
// ---- Invoke callback for first token ----
|
| 897 |
+
const char *token_text = decode_token(ctx->vocab, next_token);
|
| 898 |
+
if (callback) {
|
| 899 |
+
double gen_time = now_ms() - t0 - ctx->ttft_ms;
|
| 900 |
+
double tps = gen_time > 0 ? 1000.0 / gen_time : 0;
|
| 901 |
+
int stop = callback(token_text, next_token, ctx->tokens_generated, tps, user_data);
|
| 902 |
+
if (stop) {
|
| 903 |
+
free(pt->ids); free(pt);
|
| 904 |
+
ctx->total_time_ms = now_ms() - t0;
|
| 905 |
+
return ctx->tokens_generated;
|
| 906 |
+
}
|
| 907 |
+
}
|
| 908 |
+
|
| 909 |
+
int in_think = (next_token == cfg.think_start_token) ? 1 : 0;
|
| 910 |
+
int think_tokens = 0;
|
| 911 |
+
|
| 912 |
+
// ---- Auto-regressive generation loop ----
|
| 913 |
+
double gen_start = now_ms();
|
| 914 |
+
|
| 915 |
+
for (int gen = 1; gen < max_tokens; gen++) { @autoreleasepool {
|
| 916 |
+
// Check cancellation
|
| 917 |
+
if (atomic_load(&ctx->cancelled)) break;
|
| 918 |
+
|
| 919 |
+
// Check EOS
|
| 920 |
+
int is_eos = 0;
|
| 921 |
+
for (int e = 0; e < cfg.num_eos_tokens; e++) {
|
| 922 |
+
if (next_token == cfg.eos_token_ids[e]) { is_eos = 1; break; }
|
| 923 |
+
}
|
| 924 |
+
if (is_eos) break;
|
| 925 |
+
|
| 926 |
+
// Think budget enforcement
|
| 927 |
+
if (next_token == cfg.think_start_token) in_think = 1;
|
| 928 |
+
if (next_token == cfg.think_end_token) in_think = 0;
|
| 929 |
+
if (in_think) think_tokens++;
|
| 930 |
+
|
| 931 |
+
// Embed + forward pass
|
| 932 |
+
embed_lookup(ctx->wf, next_token, ctx->hidden);
|
| 933 |
+
|
| 934 |
+
for (int layer = 0; layer < cfg.num_layers; layer++) {
|
| 935 |
+
int is_full = cfg.is_full_attn[layer];
|
| 936 |
+
fused_layer_forward(ctx->wf, layer, ctx->hidden,
|
| 937 |
+
is_full ? ctx->kv_caches[layer] : NULL,
|
| 938 |
+
is_full ? NULL : ctx->layer_states[layer],
|
| 939 |
+
pos,
|
| 940 |
+
ctx->layer_mmaps[layer] != MAP_FAILED ? ctx->layer_mmaps[layer] : NULL,
|
| 941 |
+
K, ctx->layer_fds[layer]);
|
| 942 |
+
}
|
| 943 |
+
complete_deferred_experts();
|
| 944 |
+
pos++;
|
| 945 |
+
|
| 946 |
+
// Final norm + LM head
|
| 947 |
+
if (ctx->final_norm_w) {
|
| 948 |
+
float *normed = malloc(cfg.hidden_dim * sizeof(float));
|
| 949 |
+
cpu_rms_norm(ctx->hidden, ctx->final_norm_w, normed, cfg.hidden_dim, cfg.rms_norm_eps);
|
| 950 |
+
memcpy(ctx->hidden, normed, cfg.hidden_dim * sizeof(float));
|
| 951 |
+
free(normed);
|
| 952 |
+
}
|
| 953 |
+
|
| 954 |
+
lm_head_forward(ctx->wf, ctx->hidden, ctx->logits);
|
| 955 |
+
next_token = cpu_argmax(ctx->logits, cfg.vocab_size);
|
| 956 |
+
|
| 957 |
+
// Think budget: -1 = disabled (immediate end), >0 = max tokens, 0 = unlimited
|
| 958 |
+
if (in_think && g_think_budget < 0) {
|
| 959 |
+
next_token = cfg.think_end_token;
|
| 960 |
+
in_think = 0;
|
| 961 |
+
} else if (in_think && g_think_budget > 0 && think_tokens >= g_think_budget) {
|
| 962 |
+
next_token = cfg.think_end_token;
|
| 963 |
+
in_think = 0;
|
| 964 |
+
}
|
| 965 |
+
|
| 966 |
+
ctx->tokens_generated++;
|
| 967 |
+
|
| 968 |
+
// Compute tok/s
|
| 969 |
+
double elapsed_gen = now_ms() - gen_start;
|
| 970 |
+
ctx->tokens_per_second = elapsed_gen > 0 ? (ctx->tokens_generated - 1) * 1000.0 / elapsed_gen : 0;
|
| 971 |
+
|
| 972 |
+
// Invoke callback
|
| 973 |
+
token_text = decode_token(ctx->vocab, next_token);
|
| 974 |
+
if (callback) {
|
| 975 |
+
int stop = callback(token_text, next_token, ctx->tokens_generated,
|
| 976 |
+
ctx->tokens_per_second, user_data);
|
| 977 |
+
if (stop) break;
|
| 978 |
+
}
|
| 979 |
+
} /* @autoreleasepool */ }
|
| 980 |
+
|
| 981 |
+
ctx->total_time_ms = now_ms() - t0;
|
| 982 |
+
double gen_elapsed = now_ms() - gen_start;
|
| 983 |
+
if (ctx->tokens_generated > 1 && gen_elapsed > 0) {
|
| 984 |
+
ctx->tokens_per_second = (ctx->tokens_generated - 1) * 1000.0 / gen_elapsed;
|
| 985 |
+
}
|
| 986 |
+
|
| 987 |
+
// Persist state for KV cache reuse in next turn
|
| 988 |
+
ctx->current_pos = pos;
|
| 989 |
+
ctx->turn_count++;
|
| 990 |
+
|
| 991 |
+
free(pt->ids);
|
| 992 |
+
free(pt);
|
| 993 |
+
|
| 994 |
+
return ctx->tokens_generated;
|
| 995 |
+
}
|
| 996 |
+
}
|
| 997 |
+
|
| 998 |
+
// ============================================================================
|
| 999 |
+
// Continuation generation — reuses KV cache from previous turns
|
| 1000 |
+
// ============================================================================
|
| 1001 |
+
|
| 1002 |
+
int flashmoe_generate_continuation(
|
| 1003 |
+
FlashMoEContext *ctx,
|
| 1004 |
+
const char *user_content,
|
| 1005 |
+
int max_tokens,
|
| 1006 |
+
FlashMoETokenCallback callback,
|
| 1007 |
+
void *user_data
|
| 1008 |
+
) {
|
| 1009 |
+
if (!ctx || !ctx->loaded || !user_content) {
|
| 1010 |
+
if (ctx) snprintf(ctx->last_error, sizeof(ctx->last_error), "Engine not loaded or invalid arguments");
|
| 1011 |
+
return -1;
|
| 1012 |
+
}
|
| 1013 |
+
if (ctx->turn_count == 0) {
|
| 1014 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "No previous turn — use flashmoe_generate first");
|
| 1015 |
+
return -1;
|
| 1016 |
+
}
|
| 1017 |
+
|
| 1018 |
+
@autoreleasepool {
|
| 1019 |
+
atomic_store(&ctx->cancelled, 0);
|
| 1020 |
+
ctx->tokens_generated = 0;
|
| 1021 |
+
ctx->tokens_per_second = 0;
|
| 1022 |
+
|
| 1023 |
+
// Pre-flight memory check (iOS)
|
| 1024 |
+
#if TARGET_OS_IPHONE
|
| 1025 |
+
{
|
| 1026 |
+
size_t avail = os_proc_available_memory();
|
| 1027 |
+
if (avail < 500 * 1024 * 1024) {
|
| 1028 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error),
|
| 1029 |
+
"Insufficient memory (%.0f MB available, need 500+ MB)", avail / 1e6);
|
| 1030 |
+
return -1;
|
| 1031 |
+
}
|
| 1032 |
+
}
|
| 1033 |
+
#endif
|
| 1034 |
+
|
| 1035 |
+
double t0 = now_ms();
|
| 1036 |
+
|
| 1037 |
+
// Tokenize only the new turn (with continuation markers)
|
| 1038 |
+
PromptTokens *pt = tokenize_continuation_turn_shared(user_content);
|
| 1039 |
+
if (!pt) {
|
| 1040 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Failed to tokenize continuation turn");
|
| 1041 |
+
return -1;
|
| 1042 |
+
}
|
| 1043 |
+
|
| 1044 |
+
int K = ctx->K;
|
| 1045 |
+
int pos = ctx->current_pos; // Resume from where we left off
|
| 1046 |
+
|
| 1047 |
+
// Check we have room in the KV cache
|
| 1048 |
+
if (pos + pt->count + max_tokens > cfg.max_seq_len) {
|
| 1049 |
+
NSLog(@"[FlashMoE] Context full (%d + %d + %d > %d), resetting to fresh generation",
|
| 1050 |
+
pos, pt->count, max_tokens, cfg.max_seq_len);
|
| 1051 |
+
free(pt->ids); free(pt);
|
| 1052 |
+
// Fall back to full generation with chat template
|
| 1053 |
+
// Caller should handle this by using flashmoe_generate instead
|
| 1054 |
+
snprintf(ctx->last_error, sizeof(ctx->last_error), "Context window full, reset required");
|
| 1055 |
+
return -2; // Signal to caller: context full, need reset
|
| 1056 |
+
}
|
| 1057 |
+
|
| 1058 |
+
// NOTE: No reset_delta_net_state() — reuse KV caches and linear attention state
|
| 1059 |
+
|
| 1060 |
+
// ---- Prefill continuation tokens ----
|
| 1061 |
+
float *embed_batch = NULL;
|
| 1062 |
+
if (pt->count > 1) {
|
| 1063 |
+
embed_batch = malloc((size_t)pt->count * cfg.hidden_dim * sizeof(float));
|
| 1064 |
+
for (int i = 0; i < pt->count; i++) {
|
| 1065 |
+
embed_lookup(ctx->wf, pt->ids[i], embed_batch + (size_t)i * cfg.hidden_dim);
|
| 1066 |
+
}
|
| 1067 |
+
}
|
| 1068 |
+
|
| 1069 |
+
if (pt->count > 1) {
|
| 1070 |
+
int prefill_count = pt->count - 1;
|
| 1071 |
+
int batched_done = 0;
|
| 1072 |
+
|
| 1073 |
+
if (g_prefill_batch > 1 && embed_batch) {
|
| 1074 |
+
batched_done = batched_prefill(ctx->wf, embed_batch, prefill_count, pos,
|
| 1075 |
+
K, ctx->layer_fds, ctx->layer_mmaps,
|
| 1076 |
+
ctx->kv_caches, ctx->layer_states);
|
| 1077 |
+
if (batched_done > 0) {
|
| 1078 |
+
pos += batched_done;
|
| 1079 |
+
NSLog(@"[FlashMoE] Batched continuation prefill: %d tokens", batched_done);
|
| 1080 |
+
}
|
| 1081 |
+
}
|
| 1082 |
+
|
| 1083 |
+
for (int token_idx = batched_done; token_idx < prefill_count; token_idx++) { @autoreleasepool {
|
| 1084 |
+
if (atomic_load(&ctx->cancelled)) {
|
| 1085 |
+
free(embed_batch);
|
| 1086 |
+
free(pt->ids); free(pt);
|
| 1087 |
+
return ctx->tokens_generated;
|
| 1088 |
+
}
|
| 1089 |
+
|
| 1090 |
+
memcpy(ctx->hidden, embed_batch + (size_t)token_idx * cfg.hidden_dim,
|
| 1091 |
+
cfg.hidden_dim * sizeof(float));
|
| 1092 |
+
|
| 1093 |
+
for (int layer = 0; layer < cfg.num_layers; layer++) {
|
| 1094 |
+
int is_full = cfg.is_full_attn[layer];
|
| 1095 |
+
fused_layer_forward(ctx->wf, layer, ctx->hidden,
|
| 1096 |
+
is_full ? ctx->kv_caches[layer] : NULL,
|
| 1097 |
+
is_full ? NULL : ctx->layer_states[layer],
|
| 1098 |
+
pos,
|
| 1099 |
+
ctx->layer_mmaps[layer] != MAP_FAILED ? ctx->layer_mmaps[layer] : NULL,
|
| 1100 |
+
K, ctx->layer_fds[layer]);
|
| 1101 |
+
}
|
| 1102 |
+
discard_deferred_experts();
|
| 1103 |
+
pos++;
|
| 1104 |
+
} /* @autoreleasepool */ }
|
| 1105 |
+
}
|
| 1106 |
+
|
| 1107 |
+
// Last prefill token
|
| 1108 |
+
{
|
| 1109 |
+
if (embed_batch) {
|
| 1110 |
+
memcpy(ctx->hidden, embed_batch + (size_t)(pt->count - 1) * cfg.hidden_dim,
|
| 1111 |
+
cfg.hidden_dim * sizeof(float));
|
| 1112 |
+
} else {
|
| 1113 |
+
embed_lookup(ctx->wf, pt->ids[0], ctx->hidden);
|
| 1114 |
+
}
|
| 1115 |
+
|
| 1116 |
+
for (int layer = 0; layer < cfg.num_layers; layer++) {
|
| 1117 |
+
int is_full = cfg.is_full_attn[layer];
|
| 1118 |
+
fused_layer_forward(ctx->wf, layer, ctx->hidden,
|
| 1119 |
+
is_full ? ctx->kv_caches[layer] : NULL,
|
| 1120 |
+
is_full ? NULL : ctx->layer_states[layer],
|
| 1121 |
+
pos,
|
| 1122 |
+
ctx->layer_mmaps[layer] != MAP_FAILED ? ctx->layer_mmaps[layer] : NULL,
|
| 1123 |
+
K, ctx->layer_fds[layer]);
|
| 1124 |
+
}
|
| 1125 |
+
complete_deferred_experts();
|
| 1126 |
+
pos++;
|
| 1127 |
+
}
|
| 1128 |
+
|
| 1129 |
+
if (embed_batch) { free(embed_batch); embed_batch = NULL; }
|
| 1130 |
+
|
| 1131 |
+
// ---- Final norm + LM head + sample first token ----
|
| 1132 |
+
if (ctx->final_norm_w) {
|
| 1133 |
+
float *normed = malloc(cfg.hidden_dim * sizeof(float));
|
| 1134 |
+
cpu_rms_norm(ctx->hidden, ctx->final_norm_w, normed, cfg.hidden_dim, cfg.rms_norm_eps);
|
| 1135 |
+
memcpy(ctx->hidden, normed, cfg.hidden_dim * sizeof(float));
|
| 1136 |
+
free(normed);
|
| 1137 |
+
}
|
| 1138 |
+
|
| 1139 |
+
lm_head_forward(ctx->wf, ctx->hidden, ctx->logits);
|
| 1140 |
+
int next_token = cpu_argmax(ctx->logits, cfg.vocab_size);
|
| 1141 |
+
|
| 1142 |
+
ctx->ttft_ms = now_ms() - t0;
|
| 1143 |
+
ctx->tokens_generated = 1;
|
| 1144 |
+
|
| 1145 |
+
const char *token_text = decode_token(ctx->vocab, next_token);
|
| 1146 |
+
if (callback) {
|
| 1147 |
+
double gen_time = now_ms() - t0 - ctx->ttft_ms;
|
| 1148 |
+
double tps = gen_time > 0 ? 1000.0 / gen_time : 0;
|
| 1149 |
+
int stop = callback(token_text, next_token, ctx->tokens_generated, tps, user_data);
|
| 1150 |
+
if (stop) {
|
| 1151 |
+
free(pt->ids); free(pt);
|
| 1152 |
+
ctx->current_pos = pos;
|
| 1153 |
+
ctx->total_time_ms = now_ms() - t0;
|
| 1154 |
+
return ctx->tokens_generated;
|
| 1155 |
+
}
|
| 1156 |
+
}
|
| 1157 |
+
|
| 1158 |
+
int in_think = (next_token == cfg.think_start_token) ? 1 : 0;
|
| 1159 |
+
int think_tokens = 0;
|
| 1160 |
+
|
| 1161 |
+
// ---- Auto-regressive generation loop ----
|
| 1162 |
+
double gen_start = now_ms();
|
| 1163 |
+
|
| 1164 |
+
for (int gen = 1; gen < max_tokens; gen++) { @autoreleasepool {
|
| 1165 |
+
if (atomic_load(&ctx->cancelled)) break;
|
| 1166 |
+
|
| 1167 |
+
int is_eos = 0;
|
| 1168 |
+
for (int e = 0; e < cfg.num_eos_tokens; e++) {
|
| 1169 |
+
if (next_token == cfg.eos_token_ids[e]) { is_eos = 1; break; }
|
| 1170 |
+
}
|
| 1171 |
+
if (is_eos) break;
|
| 1172 |
+
|
| 1173 |
+
if (next_token == cfg.think_start_token) in_think = 1;
|
| 1174 |
+
if (next_token == cfg.think_end_token) in_think = 0;
|
| 1175 |
+
if (in_think) think_tokens++;
|
| 1176 |
+
|
| 1177 |
+
embed_lookup(ctx->wf, next_token, ctx->hidden);
|
| 1178 |
+
|
| 1179 |
+
for (int layer = 0; layer < cfg.num_layers; layer++) {
|
| 1180 |
+
int is_full = cfg.is_full_attn[layer];
|
| 1181 |
+
fused_layer_forward(ctx->wf, layer, ctx->hidden,
|
| 1182 |
+
is_full ? ctx->kv_caches[layer] : NULL,
|
| 1183 |
+
is_full ? NULL : ctx->layer_states[layer],
|
| 1184 |
+
pos,
|
| 1185 |
+
ctx->layer_mmaps[layer] != MAP_FAILED ? ctx->layer_mmaps[layer] : NULL,
|
| 1186 |
+
K, ctx->layer_fds[layer]);
|
| 1187 |
+
}
|
| 1188 |
+
complete_deferred_experts();
|
| 1189 |
+
pos++;
|
| 1190 |
+
|
| 1191 |
+
if (ctx->final_norm_w) {
|
| 1192 |
+
float *normed = malloc(cfg.hidden_dim * sizeof(float));
|
| 1193 |
+
cpu_rms_norm(ctx->hidden, ctx->final_norm_w, normed, cfg.hidden_dim, cfg.rms_norm_eps);
|
| 1194 |
+
memcpy(ctx->hidden, normed, cfg.hidden_dim * sizeof(float));
|
| 1195 |
+
free(normed);
|
| 1196 |
+
}
|
| 1197 |
+
|
| 1198 |
+
lm_head_forward(ctx->wf, ctx->hidden, ctx->logits);
|
| 1199 |
+
next_token = cpu_argmax(ctx->logits, cfg.vocab_size);
|
| 1200 |
+
|
| 1201 |
+
if (in_think && g_think_budget > 0 && think_tokens >= g_think_budget) {
|
| 1202 |
+
next_token = cfg.think_end_token;
|
| 1203 |
+
in_think = 0;
|
| 1204 |
+
}
|
| 1205 |
+
|
| 1206 |
+
ctx->tokens_generated++;
|
| 1207 |
+
double elapsed_gen = now_ms() - gen_start;
|
| 1208 |
+
ctx->tokens_per_second = elapsed_gen > 0 ? (ctx->tokens_generated - 1) * 1000.0 / elapsed_gen : 0;
|
| 1209 |
+
|
| 1210 |
+
token_text = decode_token(ctx->vocab, next_token);
|
| 1211 |
+
if (callback) {
|
| 1212 |
+
int stop = callback(token_text, next_token, ctx->tokens_generated,
|
| 1213 |
+
ctx->tokens_per_second, user_data);
|
| 1214 |
+
if (stop) break;
|
| 1215 |
+
}
|
| 1216 |
+
} /* @autoreleasepool */ }
|
| 1217 |
+
|
| 1218 |
+
ctx->total_time_ms = now_ms() - t0;
|
| 1219 |
+
double gen_elapsed = now_ms() - gen_start;
|
| 1220 |
+
if (ctx->tokens_generated > 1 && gen_elapsed > 0) {
|
| 1221 |
+
ctx->tokens_per_second = (ctx->tokens_generated - 1) * 1000.0 / gen_elapsed;
|
| 1222 |
+
}
|
| 1223 |
+
|
| 1224 |
+
ctx->current_pos = pos;
|
| 1225 |
+
ctx->turn_count++;
|
| 1226 |
+
|
| 1227 |
+
free(pt->ids);
|
| 1228 |
+
free(pt);
|
| 1229 |
+
|
| 1230 |
+
return ctx->tokens_generated;
|
| 1231 |
+
}
|
| 1232 |
+
}
|
| 1233 |
+
|
| 1234 |
+
void flashmoe_cancel(FlashMoEContext *ctx) {
|
| 1235 |
+
if (!ctx) return;
|
| 1236 |
+
atomic_store(&ctx->cancelled, 1);
|
| 1237 |
+
}
|
| 1238 |
+
|
| 1239 |
+
void flashmoe_reset(FlashMoEContext *ctx) {
|
| 1240 |
+
if (!ctx || !ctx->loaded) return;
|
| 1241 |
+
|
| 1242 |
+
@autoreleasepool {
|
| 1243 |
+
// Wait for any in-flight GPU work
|
| 1244 |
+
if (g_deferred.active) {
|
| 1245 |
+
[g_deferred.cmd_experts waitUntilCompleted];
|
| 1246 |
+
g_deferred.active = 0;
|
| 1247 |
+
g_deferred.cmd_experts = nil;
|
| 1248 |
+
}
|
| 1249 |
+
|
| 1250 |
+
// Reset delta-net state
|
| 1251 |
+
reset_delta_net_state();
|
| 1252 |
+
|
| 1253 |
+
// Reset KV caches
|
| 1254 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 1255 |
+
if (ctx->kv_caches[i]) {
|
| 1256 |
+
ctx->kv_caches[i]->len = 0;
|
| 1257 |
+
}
|
| 1258 |
+
}
|
| 1259 |
+
|
| 1260 |
+
// Reset conversation position
|
| 1261 |
+
ctx->current_pos = 0;
|
| 1262 |
+
ctx->turn_count = 0;
|
| 1263 |
+
|
| 1264 |
+
// Reset stats
|
| 1265 |
+
ctx->tokens_generated = 0;
|
| 1266 |
+
ctx->tokens_per_second = 0;
|
| 1267 |
+
ctx->total_time_ms = 0;
|
| 1268 |
+
ctx->ttft_ms = 0;
|
| 1269 |
+
}
|
| 1270 |
+
}
|
| 1271 |
+
|
| 1272 |
+
void flashmoe_get_stats(FlashMoEContext *ctx, FlashMoEStats *stats) {
|
| 1273 |
+
if (!ctx || !stats) return;
|
| 1274 |
+
|
| 1275 |
+
memset(stats, 0, sizeof(FlashMoEStats));
|
| 1276 |
+
|
| 1277 |
+
if (ctx->loaded) {
|
| 1278 |
+
snprintf(stats->model_name, sizeof(stats->model_name), "%s", cfg.model_path);
|
| 1279 |
+
stats->num_layers = cfg.num_layers;
|
| 1280 |
+
stats->num_experts = cfg.num_experts;
|
| 1281 |
+
stats->active_experts_k = ctx->K;
|
| 1282 |
+
stats->hidden_dim = cfg.hidden_dim;
|
| 1283 |
+
stats->vocab_size = cfg.vocab_size;
|
| 1284 |
+
stats->weight_file_bytes = ctx->wf ? ctx->wf->size : 0;
|
| 1285 |
+
|
| 1286 |
+
// Compute total expert file bytes
|
| 1287 |
+
size_t total_expert = 0;
|
| 1288 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 1289 |
+
total_expert += ctx->layer_mmap_sizes[i];
|
| 1290 |
+
}
|
| 1291 |
+
stats->expert_file_bytes = total_expert;
|
| 1292 |
+
|
| 1293 |
+
// Approximate Metal buffer bytes
|
| 1294 |
+
stats->metal_buffer_bytes = (size_t)cfg.expert_size_4bit * MAX_K * 2 + // expert data (double-buffered)
|
| 1295 |
+
(size_t)cfg.hidden_dim * sizeof(float) * 20 + // various working buffers
|
| 1296 |
+
(size_t)cfg.vocab_size * sizeof(float); // logits
|
| 1297 |
+
}
|
| 1298 |
+
|
| 1299 |
+
stats->tokens_per_second = ctx->tokens_per_second;
|
| 1300 |
+
stats->tokens_generated = ctx->tokens_generated;
|
| 1301 |
+
stats->total_time_ms = ctx->total_time_ms;
|
| 1302 |
+
stats->ttft_ms = ctx->ttft_ms;
|
| 1303 |
+
}
|
| 1304 |
+
|
| 1305 |
+
void flashmoe_set_runtime_config(FlashMoEContext *ctx,
|
| 1306 |
+
int active_experts_k,
|
| 1307 |
+
int cmd_merge,
|
| 1308 |
+
int fused_attention,
|
| 1309 |
+
int cache_io_split,
|
| 1310 |
+
int fp16_accumulation) {
|
| 1311 |
+
if (!ctx || !ctx->loaded) return;
|
| 1312 |
+
|
| 1313 |
+
// K: clamp to [1, model default]
|
| 1314 |
+
int k = active_experts_k;
|
| 1315 |
+
if (k <= 0) k = cfg.num_experts_per_tok;
|
| 1316 |
+
if (k > cfg.num_experts_per_tok) k = cfg.num_experts_per_tok;
|
| 1317 |
+
if (k > MAX_K) k = MAX_K;
|
| 1318 |
+
ctx->K = k;
|
| 1319 |
+
|
| 1320 |
+
// Set C globals directly — no reload needed
|
| 1321 |
+
g_cmd_merge_enabled = cmd_merge ? 1 : 0;
|
| 1322 |
+
g_fused_attention_enabled = fused_attention ? 1 : 0;
|
| 1323 |
+
g_cache_io_split = cache_io_split > 1 ? cache_io_split : 1;
|
| 1324 |
+
g_use_fp16_accum = fp16_accumulation ? 1 : 0;
|
| 1325 |
+
}
|
| 1326 |
+
|
| 1327 |
+
int flashmoe_validate_model(const char *model_path) {
|
| 1328 |
+
if (!model_path) return -1;
|
| 1329 |
+
|
| 1330 |
+
// Check config.json
|
| 1331 |
+
char path[1024];
|
| 1332 |
+
snprintf(path, sizeof(path), "%s/config.json", model_path);
|
| 1333 |
+
if (access(path, R_OK) != 0) return -1;
|
| 1334 |
+
|
| 1335 |
+
// Check model_weights.bin (or split: model_weights_0.bin + model_weights_1.bin)
|
| 1336 |
+
snprintf(path, sizeof(path), "%s/model_weights.bin", model_path);
|
| 1337 |
+
if (access(path, R_OK) != 0) {
|
| 1338 |
+
char path0[1024], path1[1024];
|
| 1339 |
+
snprintf(path0, sizeof(path0), "%s/model_weights_0.bin", model_path);
|
| 1340 |
+
snprintf(path1, sizeof(path1), "%s/model_weights_1.bin", model_path);
|
| 1341 |
+
if (access(path0, R_OK) != 0 || access(path1, R_OK) != 0) return -1;
|
| 1342 |
+
}
|
| 1343 |
+
|
| 1344 |
+
// Check model_weights.json
|
| 1345 |
+
snprintf(path, sizeof(path), "%s/model_weights.json", model_path);
|
| 1346 |
+
if (access(path, R_OK) != 0) return -1;
|
| 1347 |
+
|
| 1348 |
+
// Check for at least one expert layer file
|
| 1349 |
+
snprintf(path, sizeof(path), "%s/packed_experts/layer_00.bin", model_path);
|
| 1350 |
+
int has_4bit = (access(path, R_OK) == 0);
|
| 1351 |
+
|
| 1352 |
+
snprintf(path, sizeof(path), "%s/packed_experts_tiered/layer_00.bin", model_path);
|
| 1353 |
+
int has_tiered = (access(path, R_OK) == 0);
|
| 1354 |
+
|
| 1355 |
+
snprintf(path, sizeof(path), "%s/packed_experts_2bit/layer_00.bin", model_path);
|
| 1356 |
+
int has_2bit = (access(path, R_OK) == 0);
|
| 1357 |
+
|
| 1358 |
+
if (!has_4bit && !has_tiered && !has_2bit) return -1;
|
| 1359 |
+
|
| 1360 |
+
return 0;
|
| 1361 |
+
}
|
| 1362 |
+
|
| 1363 |
+
int flashmoe_turn_count(FlashMoEContext *ctx) {
|
| 1364 |
+
if (!ctx) return 0;
|
| 1365 |
+
return ctx->turn_count;
|
| 1366 |
+
}
|
| 1367 |
+
|
| 1368 |
+
const char *flashmoe_last_error(FlashMoEContext *ctx) {
|
| 1369 |
+
if (!ctx) return "NULL context";
|
| 1370 |
+
return ctx->last_error;
|
| 1371 |
+
}
|
| 1372 |
+
|
FlashMoE-iOS/IOS_PORT.md
ADDED
|
@@ -0,0 +1,347 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Flash-MoE iOS Port: From MacBook to iPhone
|
| 2 |
+
|
| 3 |
+
How we took a 397B-parameter MoE inference engine written in C/Metal and made it run on an iPhone — getting 5.5 tok/s on the 35B model, with the 397B model loading but hitting hard platform limits.
|
| 4 |
+
|
| 5 |
+
## The Challenge
|
| 6 |
+
|
| 7 |
+
The desktop Flash-MoE engine was designed for a MacBook Pro with 48GB unified memory, 40-core GPU, and 17.5 GB/s NVMe. iPhones have 8-12GB RAM, ~10-core GPU, ~2.5 GB/s NVMe, and no filesystem access to `shaders.metal` at runtime. Every assumption in the 7,500-line inference engine needed to be re-examined.
|
| 8 |
+
|
| 9 |
+
## Results
|
| 10 |
+
|
| 11 |
+
| Device | Model | K | tok/s | Notes |
|
| 12 |
+
|--------|-------|---|-------|-------|
|
| 13 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-35B-A3B | 8 | **5.5** | Full quality, 19.5GB download, full GPU path |
|
| 14 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-35B-A3B (tiered) | 8 | **5.5+** | 13.4GB download, same quality |
|
| 15 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-397B-A17B | 4 | ~0.003* | *CPU fallback only. 6 min/token. |
|
| 16 |
+
| MacBook Pro M3 Max (48GB) | Qwen3.5-35B-A3B | 8 | **9.7** | After autoresearch optimizations |
|
| 17 |
+
| MacBook Pro M3 Max (48GB) | Qwen3.5-397B-A17B | 4 | **4.4** | K-reduced from K=10 |
|
| 18 |
+
|
| 19 |
+
*397B model loads and generates correct tokens but uses CPU fallback for weight matmuls because the 5.5GB `model_weights.bin` cannot fit in Metal buffer(s) on a 12GB device. See [397B_ANALYSIS.md](397B_ANALYSIS.md) for the full breakdown.
|
| 20 |
+
|
| 21 |
+
iPhone achieves **57% of laptop speed** on the 35B model with 17% of the memory.
|
| 22 |
+
|
| 23 |
+
## What We Built
|
| 24 |
+
|
| 25 |
+
A native SwiftUI iOS app wrapping the C/Metal inference engine with:
|
| 26 |
+
|
| 27 |
+
- **Model discovery**: on-device scanning + HuggingFace download catalog
|
| 28 |
+
- **Interactive chat**: streaming tokens, thinking animation with disclosure triangle, text selection
|
| 29 |
+
- **Model management**: import from Files app, export via `moveToService`, swipe-to-delete
|
| 30 |
+
- **Settings**: K-reduction picker (K=2/4/6/8/10), I/O fanout picker (off/2/4/8 chunks)
|
| 31 |
+
- **Profiler view**: resource monitoring overlay with thermal state indicator (Cool/Warm/Hot/Critical)
|
| 32 |
+
- **Multi-turn chat**: KV cache reuse across chat turns, adaptive context length by device RAM
|
| 33 |
+
- **Background downloads**: `URLSession` downloads from HuggingFace with resume support
|
| 34 |
+
- **Cross-app model access**: security-scoped bookmarks for reading models from other app containers
|
| 35 |
+
|
| 36 |
+
**Architecture**: SwiftUI (UI + @Observable state) -> Swift async bridge (AsyncStream) -> Objective-C wrapper -> C inference engine (7,500 lines) -> Metal GPU shaders (1,200 lines)
|
| 37 |
+
|
| 38 |
+
## Performance Optimizations (Autoresearch)
|
| 39 |
+
|
| 40 |
+
Before the iOS port, we ran Karpathy's autoresearch pattern — an autonomous experiment loop that modifies Metal shaders, benchmarks, and keeps/discards based on tok/s. 10 experiments, 4 kept:
|
| 41 |
+
|
| 42 |
+
| Experiment | Description | Impact | Why It Works |
|
| 43 |
+
|-----------|-------------|--------|-------------|
|
| 44 |
+
| SIMD reduction | Replace serial thread-0 accumulation in `rms_norm_qk`/`gated_rms_norm` with `simd_sum` + shared memory reduction | +2.1% | Eliminates serial bottleneck in GPU reduction — 4 SIMD groups of 32 threads each contribute partial sums |
|
| 45 |
+
| FMA 2-bit kernel | Apply `fma(nibble, scale*x, bias*x)` pattern to 2-bit dequant (same trick as 4-bit v3) | +6.2% | GPU fused multiply-add does dequant+multiply in one instruction instead of two |
|
| 46 |
+
| Half-precision x_shared (v3) | Store threadgroup shared memory input cache as `half` instead of `float` | +12.1% | Halves shared memory from 16KB to 8KB, doubles GPU occupancy. Input values are already approximate from dequantization so half precision loses nothing |
|
| 47 |
+
| Half-precision x_shared (2-bit) | Same trick applied to 2-bit kernel | +3.3% | Same occupancy benefit, smaller because 2-bit kernel is more I/O bound |
|
| 48 |
+
|
| 49 |
+
**Combined: +15.3% theoretical, +34.7% real-world** (from 7.2 to 9.7 tok/s on MacBook Pro).
|
| 50 |
+
|
| 51 |
+
6 experiments discarded (single compute encoder -5%, FMA without shared memory -15%, extended x_shared to 8192 -7%, three others marginal).
|
| 52 |
+
|
| 53 |
+
## Problems Solved (iOS Port)
|
| 54 |
+
|
| 55 |
+
### 1. Metal Shader Loading
|
| 56 |
+
|
| 57 |
+
**Problem**: The desktop engine compiles `shaders.metal` from source at runtime via `newLibraryWithSource:`. On iOS, there is no filesystem path to the shader file.
|
| 58 |
+
|
| 59 |
+
**Fix**: Runtime fallback — try `[device newDefaultLibrary]` first (loads pre-compiled `default.metallib` from the app bundle), fall back to source compilation for macOS CLI:
|
| 60 |
+
|
| 61 |
+
```objc
|
| 62 |
+
ctx->library = [ctx->device newDefaultLibrary];
|
| 63 |
+
if (ctx->library) {
|
| 64 |
+
// iOS: loaded from bundle
|
| 65 |
+
} else {
|
| 66 |
+
// macOS: compile from source
|
| 67 |
+
NSString *src = [NSString stringWithContentsOfFile:@"shaders.metal" ...];
|
| 68 |
+
ctx->library = [ctx->device newLibraryWithSource:src ...];
|
| 69 |
+
}
|
| 70 |
+
```
|
| 71 |
+
|
| 72 |
+
**Xcode fix**: Moved `shaders.metal` from the **Resources** build phase to **Sources** so Xcode's Metal compiler produces `default.metallib` in the app bundle.
|
| 73 |
+
|
| 74 |
+
### 2. Memory: KV Cache OOM
|
| 75 |
+
|
| 76 |
+
**Problem**: The model's `max_position_embeddings` is 131,072 (128k context). KV cache allocation per full-attention layer: `131072 * kv_heads * head_dim * 4 bytes`. For the 35B (10 full-attn layers): ~2.5GB just for KV caches. `calloc` silently returns NULL on iPhone, causing `EXC_BAD_ACCESS` crashes.
|
| 77 |
+
|
| 78 |
+
**Fix**: Adaptive context length based on `os_proc_available_memory()`. Budget 25% of available memory for KV caches, clamp to power-of-2 sizes (512-8192). On a 12GB iPhone this yields ~2048 context, keeping KV caches under ~40MB total. Context window is limited but sufficient for chat.
|
| 79 |
+
|
| 80 |
+
### 3. Memory: Metal Debug Layer
|
| 81 |
+
|
| 82 |
+
**Problem**: Debug builds wrap every Metal object with `MTLDebugComputeCommandEncoder` validation proxies, adding ~2GB of overhead — roughly doubling GPU memory usage. The app crashes with `NSMallocException` trying to allocate debug wrappers.
|
| 83 |
+
|
| 84 |
+
**Fix**: Build in **Release** mode and disable Metal API Validation in the Xcode scheme. The debug overhead is too large for iPhone's memory budget.
|
| 85 |
+
|
| 86 |
+
### 4. Tokenizer Not Found
|
| 87 |
+
|
| 88 |
+
**Problem**: `init_tokenizer()` searches for `tokenizer.bin` at relative filesystem paths (`./tokenizer.bin`, `./metal_infer/tokenizer.bin`). These don't exist on iOS.
|
| 89 |
+
|
| 90 |
+
**Fix**: Extended the search to check the model directory (where it's downloaded) and the app bundle:
|
| 91 |
+
|
| 92 |
+
```objc
|
| 93 |
+
// Try model directory (downloaded with model)
|
| 94 |
+
snprintf(model_tok, sizeof(model_tok), "%s/tokenizer.bin", cfg.model_path);
|
| 95 |
+
if (access(model_tok, R_OK) == 0) { bpe_load(&g_tokenizer, model_tok); }
|
| 96 |
+
|
| 97 |
+
// Try app bundle
|
| 98 |
+
NSString *bundlePath = [[NSBundle mainBundle] pathForResource:@"tokenizer" ofType:@"bin"];
|
| 99 |
+
if (bundlePath) { bpe_load(&g_tokenizer, [bundlePath UTF8String]); }
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
### 5. Missing Info.plist
|
| 103 |
+
|
| 104 |
+
**Problem**: Xcode target configs had `INFOPLIST_KEY_*` entries but never set `GENERATE_INFOPLIST_FILE = YES`, so no `Info.plist` was produced.
|
| 105 |
+
|
| 106 |
+
**Fix**: Added `GENERATE_INFOPLIST_FILE = YES` to both Debug and Release target build settings.
|
| 107 |
+
|
| 108 |
+
### 6. Chat Template (Garbage Output)
|
| 109 |
+
|
| 110 |
+
**Problem**: The model received raw text ("Hi") instead of Qwen's chat template format. Without the `<|im_start|>` / `<|im_end|>` markers, the model treats input as a continuation of arbitrary text, producing incoherent output.
|
| 111 |
+
|
| 112 |
+
**Fix**: Added `buildChatPrompt()` in `ChatView.swift` that formats the full conversation history using Qwen's `<|im_start|>system/user/assistant<|im_end|>` template.
|
| 113 |
+
|
| 114 |
+
### 7. Special Token Leakage
|
| 115 |
+
|
| 116 |
+
**Problem**: End-of-turn tokens like `<|im_end|>` appear as visible text in the chat UI.
|
| 117 |
+
|
| 118 |
+
**Fix**: Strip special tokens from the token stream before displaying.
|
| 119 |
+
|
| 120 |
+
### 8. KV Cache Reuse Across Turns
|
| 121 |
+
|
| 122 |
+
**Problem**: Each new message re-tokenizes and re-prefills the entire conversation history. On a model generating at 5.5 tok/s, re-prefilling 200+ tokens of history adds seconds of latency per turn.
|
| 123 |
+
|
| 124 |
+
**Fix**: Added `flashmoe_generate_continuation()` — a new C API that reuses existing KV cache state. The engine tracks conversation position and only processes the new user turn tokens. Returns -2 if context is full, signaling the Swift bridge to reset and do a full re-prefill.
|
| 125 |
+
|
| 126 |
+
### 9. MAX_K Buffer Overflow
|
| 127 |
+
|
| 128 |
+
**Problem**: The engine hardcodes `#define MAX_K 8` for multi-expert buffer arrays. The 397B model needs K=10 (default). Without K-reduction (or bumping MAX_K), loading the 397B model causes a buffer overflow crash writing past `buf_multi_expert_data[8]`.
|
| 129 |
+
|
| 130 |
+
**Fix**: Bumped MAX_K to 16 and added a runtime cap `min(K, MAX_K)` as a safety net.
|
| 131 |
+
|
| 132 |
+
### 10. Race Condition in Async pread
|
| 133 |
+
|
| 134 |
+
**Problem**: The pthread-pool-based async pread had generation counter conflicts when multiple dispatch groups accessed shared state.
|
| 135 |
+
|
| 136 |
+
**Fix**: Replaced pthread pool with GCD `dispatch_group` for expert reads. Each expert read gets its own dispatch group, eliminating the shared counter race.
|
| 137 |
+
|
| 138 |
+
### 11. Tiered Validation Bug
|
| 139 |
+
|
| 140 |
+
**Problem**: `async_pread_wait` validated each chunk against a uniform 4-bit expert size, but cold experts in tiered mode are 2-bit (44% smaller). This caused silent skipping of cold experts.
|
| 141 |
+
|
| 142 |
+
**Fix**: Validate each chunk against its own actual size (resolved from the expert's quantization tier), not the uniform 4-bit size.
|
| 143 |
+
|
| 144 |
+
## Metal/GPU on iOS: The 4GB Wall
|
| 145 |
+
|
| 146 |
+
### The Hard Limit
|
| 147 |
+
|
| 148 |
+
Metal on iOS has a **hard 4GB per-buffer limit** that cannot be lifted by any entitlement. The `extended-virtual-addressing` and `increased-memory-limit` entitlements expand total addressable memory but do NOT increase the per-buffer cap.
|
| 149 |
+
|
| 150 |
+
| Scenario | Weight Size | Fits? | Result |
|
| 151 |
+
|----------|------------|-------|--------|
|
| 152 |
+
| 35B model (4-bit) | ~2.5 GB | Yes | Single Metal buffer, full GPU path, 5.5 tok/s |
|
| 153 |
+
| 397B model (4-bit) | ~5.5 GB | No | Exceeds 4096 MB Metal buffer limit |
|
| 154 |
+
|
| 155 |
+
### What We Tried to Get 397B Weights on GPU
|
| 156 |
+
|
| 157 |
+
| Approach | Result | Why It Failed |
|
| 158 |
+
|----------|--------|--------------|
|
| 159 |
+
| **Single Metal buffer (5.5GB)** | Metal assertion crash | Buffer must not exceed 4096 MB. Hard limit. |
|
| 160 |
+
| **Two overlapping Metal buffers** | OOM kill | Metal tracks ~8GB of shared memory on a 12GB device. iOS kills the app even in Release mode. |
|
| 161 |
+
| **Staging buffer (50MB memcpy per dispatch)** | Data corruption | One staging buffer shared across multiple in-flight command buffers. Metal executes all encoders AFTER commit, so earlier tensor data gets overwritten by later memcpys before GPU reads it. **Fundamental flaw** — would need N staging buffers for N concurrent command buffers. |
|
| 162 |
+
| **CPU fallback (no Metal buffers for weights)** | Works but 6 min/token | Weight matmuls run on CPU via Accelerate. Expert forward still uses GPU (separate small buffers). Correct output but unusable speed. |
|
| 163 |
+
|
| 164 |
+
### The Path Forward: Split Weight Files
|
| 165 |
+
|
| 166 |
+
The solution is to split `model_weights.bin` at the Python packing stage into two files, each under 4GB. Each file gets its own Metal buffer. This requires:
|
| 167 |
+
|
| 168 |
+
1. Modify `extract_weights.py` to emit `model_weights_0.bin` and `model_weights_1.bin` with a split point
|
| 169 |
+
2. Update `model_weights.json` manifest with file index per tensor
|
| 170 |
+
3. Update `infer.m` to mmap both files and select the correct base pointer per tensor
|
| 171 |
+
4. Re-upload 397B model to HuggingFace with split weights
|
| 172 |
+
|
| 173 |
+
This is the next major engineering task.
|
| 174 |
+
|
| 175 |
+
## K-Reduction: Running 397B on iPhone
|
| 176 |
+
|
| 177 |
+
### The Insight
|
| 178 |
+
|
| 179 |
+
Mixture-of-Experts gives a natural inference knob that dense models don't have: activate **fewer experts per token** at inference time, even if the model was trained with more. K=4 instead of K=10 means:
|
| 180 |
+
|
| 181 |
+
- **60% less I/O per token** (4 expert reads instead of 10, per layer)
|
| 182 |
+
- Storage unchanged (all 512 experts still on disk — routing decides which 4 to use)
|
| 183 |
+
- Quality degrades gracefully — you're still selecting the *best* 4 from 512 options
|
| 184 |
+
|
| 185 |
+
### K-Reduction Quality Findings
|
| 186 |
+
|
| 187 |
+
| Model | K | Output Quality | Notes |
|
| 188 |
+
|-------|---|---------------|-------|
|
| 189 |
+
| 397B (default K=10) | K=2 | **Gibberish** | "no manager" + random tokens |
|
| 190 |
+
| 397B (default K=10) | K=4 | **Degenerate** | "!!!!" repeated output |
|
| 191 |
+
| 397B (default K=10) | K=6 | Untested | Needs GPU path (split weights) to evaluate |
|
| 192 |
+
| 397B (default K=10) | K=8/K=10 | Untested | CPU path too slow (6 min/token) to evaluate quality |
|
| 193 |
+
| 35B (default K=8) | K=8 | Excellent | Full quality, production-ready |
|
| 194 |
+
|
| 195 |
+
**Conclusion**: K-reduction quality depends heavily on the model. The 397B was trained with K=10 and may need K=6+ for coherent output. Cannot properly evaluate until split weight files enable the GPU path. The 35B works perfectly at its default K=8.
|
| 196 |
+
|
| 197 |
+
### 397B on iPhone: Memory Analysis
|
| 198 |
+
|
| 199 |
+
```
|
| 200 |
+
Non-expert weights (mmap'd): 5.5 GB (virtual, not all resident)
|
| 201 |
+
Metal buffers: ~500 MB (KV cache at 2048 ctx + delta-net + expert buffers)
|
| 202 |
+
iOS overhead: ~2 GB
|
| 203 |
+
--------
|
| 204 |
+
Resident estimate: ~3 GB
|
| 205 |
+
Available for page cache: ~9 GB on 12GB iPhone
|
| 206 |
+
|
| 207 |
+
Expert I/O per token (K=4): 4 x 60 layers x 6.75 MB = 1.6 GB
|
| 208 |
+
iPhone NVMe throughput: ~2.5 GB/s
|
| 209 |
+
I/O time per token: ~0.65s
|
| 210 |
+
GPU compute per token: ~0.3-0.5s (60 layers, 32 heads, head_dim=256)
|
| 211 |
+
--------
|
| 212 |
+
Expected (with GPU path): ~1-1.5 tok/s
|
| 213 |
+
```
|
| 214 |
+
|
| 215 |
+
With tiered experts (hot=4-bit, cold=2-bit): I/O drops ~34%, storage drops from 208GB to ~140GB, expected speed improves to ~1.5-2 tok/s.
|
| 216 |
+
|
| 217 |
+
## App Container / Bundle ID Saga
|
| 218 |
+
|
| 219 |
+
Getting models onto the device and keeping them across builds was harder than the actual porting work.
|
| 220 |
+
|
| 221 |
+
### The Problem Chain
|
| 222 |
+
|
| 223 |
+
1. Original app used `com.flashmoe.ios` bundle ID on a personal (free) developer team
|
| 224 |
+
2. Memory entitlements (`extended-virtual-addressing`, `increased-memory-limit`) require a **paid** developer account
|
| 225 |
+
3. Switched to paid team, but the original bundle ID was claimed by the personal team
|
| 226 |
+
4. Apple takes 24-48 hours to release a bundle ID after the old app is deleted from the personal team
|
| 227 |
+
5. Had to create a new bundle ID: `com.alexintosh.flashmoe`
|
| 228 |
+
6. This creates a NEW app container — the old app's Documents folder (with 300GB of model data) is inaccessible
|
| 229 |
+
|
| 230 |
+
### Moving 300GB Between App Containers
|
| 231 |
+
|
| 232 |
+
| Approach | Result | Why |
|
| 233 |
+
|----------|--------|-----|
|
| 234 |
+
| Files app direct access | Failed | `UIFileSharingEnabled` was set on old app but Files never showed its Documents. Unknown why. |
|
| 235 |
+
| `pymobiledevice3` (HouseArrest/AFC) | Can read both containers | But cannot cross-container move due to iOS sandbox. Would need download-to-Mac-then-upload round-trip. |
|
| 236 |
+
| `moveToService` picker in old app | **Worked** | `UIDocumentPickerViewController` with `.moveToService` lets iOS move files to a shared location (iCloud Drive, local "On My iPhone"). Then import picker in new app reads from that location. |
|
| 237 |
+
|
| 238 |
+
**Final workflow**: Old app exports via moveToService -> shared "On My iPhone" folder -> New app imports via document picker -> Files copied to new app's Documents.
|
| 239 |
+
|
| 240 |
+
### Developer Experience Issues
|
| 241 |
+
|
| 242 |
+
| Issue | Impact | Mitigation |
|
| 243 |
+
|-------|--------|-----------|
|
| 244 |
+
| Xcode Clean Build wipes app Documents | 300GB model data lost on every Clean Build | Avoid Clean Build. Use incremental. |
|
| 245 |
+
| Debug builds add ~2GB Metal validation overhead | OOM on iPhone for any model | Always build Release for on-device testing |
|
| 246 |
+
| USB developer disk image mount failures after OOM crashes | Cannot deploy until fixed | Restart iPhone + Xcode |
|
| 247 |
+
| "Persist container data on reinstall" unavailable on personal team | Reinstall = lose all models | Required paid developer team |
|
| 248 |
+
| `isExcludedFromBackup` needed on model files | Without it, iOS may purge 200GB+ of model data | Set on every downloaded file |
|
| 249 |
+
|
| 250 |
+
### File Provider Storage Gotcha
|
| 251 |
+
|
| 252 |
+
When a model is stored via Files app "Move to Service", the actual bytes end up in `/Shared/AppGroup/.../File Provider Storage/` which goes through the iOS file coordination layer. Every `pread` call to expert files goes through this layer, adding measurable latency per I/O. Solution: always import models into the app's own Documents directory rather than reading them in place from File Provider Storage.
|
| 253 |
+
|
| 254 |
+
## iOS App Architecture
|
| 255 |
+
|
| 256 |
+
### Engine Layer (C/Objective-C)
|
| 257 |
+
|
| 258 |
+
- **FlashMoEEngine.h** — C API: `create/load/generate/generate_continuation/cancel/reset/unload/destroy`
|
| 259 |
+
- **FlashMoEEngine.m** — Unity build wrapper (`#include "infer.m"`)
|
| 260 |
+
- **infer.m** — The full 7,500-line inference engine, shared with macOS
|
| 261 |
+
- **shaders.metal** — Metal compute kernels (~1,200 lines), compiled into `default.metallib`
|
| 262 |
+
|
| 263 |
+
The unity build means 100% of inference code is shared between iOS and macOS. No fork to maintain.
|
| 264 |
+
|
| 265 |
+
### Bridge Layer (Swift/ObjC Interop)
|
| 266 |
+
|
| 267 |
+
- **FlashMoEBridge.swift** — `@Observable` class wrapping the C API
|
| 268 |
+
- `loadModel(at:)` -> background thread -> `flashmoe_load()` with adaptive memory config
|
| 269 |
+
- `generate(prompt:)` -> `AsyncStream<GenerationToken>` via C callback bridge using `Unmanaged` pointer passing
|
| 270 |
+
- `generateContinuation(userMessage:)` -> reuses KV cache for multi-turn, returns -2 on context full
|
| 271 |
+
- State machine: `idle -> loading -> ready -> generating -> ready`
|
| 272 |
+
- Thread safety: generation runs on dedicated `DispatchQueue`, callbacks bridge to MainActor for UI
|
| 273 |
+
|
| 274 |
+
### UI Layer (SwiftUI)
|
| 275 |
+
|
| 276 |
+
- **ChatView** — Streaming chat with thinking block disclosure triangle, text selection, braille spinner animation
|
| 277 |
+
- **ModelListView** — On-device models + downloadable HuggingFace catalog with auto K-reduction recommendations
|
| 278 |
+
- **ModelDownloadRow** — Per-model download progress with pause/resume/delete
|
| 279 |
+
- **ProfilerView** — Resource monitoring: memory usage, tok/s, thermal state (Cool/Warm/Hot/Critical)
|
| 280 |
+
- **ContentView** — Root navigation with Settings (K-reduction picker, I/O fanout picker)
|
| 281 |
+
|
| 282 |
+
### Services
|
| 283 |
+
|
| 284 |
+
- **DownloadManager.swift** — Background `URLSession` with state persistence, resume support, per-file progress
|
| 285 |
+
- **ModelCatalog.swift** — Static registry of pre-packed HuggingFace repos with recommended K values per device
|
| 286 |
+
|
| 287 |
+
### Entitlements
|
| 288 |
+
|
| 289 |
+
```xml
|
| 290 |
+
<key>com.apple.developer.kernel.extended-virtual-addressing</key> <true/>
|
| 291 |
+
<key>com.apple.developer.kernel.increased-memory-limit</key> <true/>
|
| 292 |
+
```
|
| 293 |
+
|
| 294 |
+
These expand total addressable memory (needed for mmap'ing 5.5GB weight files + expert reads) but do NOT lift the 4GB per-Metal-buffer limit.
|
| 295 |
+
|
| 296 |
+
## Key Design Decisions
|
| 297 |
+
|
| 298 |
+
| Decision | Why | Alternative Considered |
|
| 299 |
+
|----------|-----|----------------------|
|
| 300 |
+
| Unity build (`#include "infer.m"`) | Share 100% of inference code with macOS, no fork | Separate iOS codebase — rejected, maintenance nightmare |
|
| 301 |
+
| Runtime Metal library fallback | Single codepath for iOS (pre-compiled) and macOS (source) | Conditional compilation — more complex |
|
| 302 |
+
| Adaptive context cap via `os_proc_available_memory()` | Scale KV cache to available memory instead of hardcoded limit | Fixed 2048 — wastes memory on larger devices |
|
| 303 |
+
| K-reduction as user setting, not hardcoded | Quality impact varies by model, let user experiment | Auto-select — not enough data on quality thresholds |
|
| 304 |
+
| Pre-packed HuggingFace models | No on-device conversion needed, download and run | On-device repacking — too slow, too much temp storage |
|
| 305 |
+
| Background URLSession | Downloads survive app suspension (not force-quit) | Foreground-only — bad UX for 200GB downloads |
|
| 306 |
+
| Trust the OS page cache on iOS | Same philosophy as desktop, no custom expert cache | Custom LRU — slower on desktop, would be worse on iOS |
|
| 307 |
+
| `isExcludedFromBackup` on model files | Prevents iOS from purging 200GB+ model data | No flag — iOS may delete models to free space |
|
| 308 |
+
| Security-scoped bookmarks for external models | Read models from other app containers without copying | Copy-on-import — wastes storage for 200GB models |
|
| 309 |
+
| `moveToService` for export | iOS handles the file move atomically | Manual copy — error-prone, needs temp space |
|
| 310 |
+
| CPU fallback for >4GB weights | Correct (just slow) vs. crashing or corrupting data | Staging buffers — data corruption, fundamental flaw |
|
| 311 |
+
| No mmap for expert files on iOS | pread-only saves virtual address space | mmap — 5x slower due to per-page fault overhead on cold data (tested on desktop) |
|
| 312 |
+
|
| 313 |
+
## Pre-Packed Models on HuggingFace
|
| 314 |
+
|
| 315 |
+
| Repo | Model | Quant | Size | iPhone Min |
|
| 316 |
+
|------|-------|-------|------|-----------|
|
| 317 |
+
| `alexintosh/Qwen3.5-35B-A3B-Q4-FlashMoE` | 35B | 4-bit | 19.5 GB | 128 GB storage, 8 GB RAM |
|
| 318 |
+
| `alexintosh/Qwen3.5-35B-A3B-Q4-Tiered-FlashMoE` | 35B | Tiered | 13.4 GB | 128 GB storage, 8 GB RAM |
|
| 319 |
+
| `alexintosh/Qwen3.5-397B-A17B-Q4-FlashMoE` | 397B | 4-bit | ~214 GB | 256 GB storage, 12 GB RAM |
|
| 320 |
+
|
| 321 |
+
Each repo contains: `config.json`, `model_weights.bin`, `model_weights.json`, `tokenizer.bin`, `tokenizer.json`, `vocab.bin`, `packed_experts/layer_XX.bin`.
|
| 322 |
+
|
| 323 |
+
## Copying Models to iPhone
|
| 324 |
+
|
| 325 |
+
```bash
|
| 326 |
+
# Via USB (auto-detects connected device)
|
| 327 |
+
./copy_model_to_iphone.sh /path/to/model-directory
|
| 328 |
+
|
| 329 |
+
# Or specify device UDID
|
| 330 |
+
./copy_model_to_iphone.sh /path/to/model-directory <device-udid>
|
| 331 |
+
```
|
| 332 |
+
|
| 333 |
+
The script copies all model files to the app's Documents container with transfer speed and ETA display.
|
| 334 |
+
|
| 335 |
+
## What's Next
|
| 336 |
+
|
| 337 |
+
### Must Do (397B on iPhone)
|
| 338 |
+
1. **Split `model_weights.bin` into two <4GB files** at Python packing stage — each gets its own Metal buffer, enabling full GPU path
|
| 339 |
+
2. **Test K=6/8/10 with GPU path** — CPU fallback is too slow to evaluate quality; need GPU path to determine minimum viable K
|
| 340 |
+
3. **Upload split 397B model to HuggingFace** — needs the split weight tooling first
|
| 341 |
+
|
| 342 |
+
### Nice to Have
|
| 343 |
+
- **Adaptive K** — auto-select K based on device RAM, thermal state, and model size
|
| 344 |
+
- **Thermal throttling awareness** — monitor `ProcessInfo.ThermalState`, reduce K when throttling
|
| 345 |
+
- **Download resumption improvements** — handle interrupted downloads more gracefully
|
| 346 |
+
- **Background inference** — continue generation when app is backgrounded
|
| 347 |
+
- **Smaller models** — 7B/14B Qwen3.5 variants for devices with less storage
|
FlashMoE-iOS/Models/ModelCatalog.swift
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* ModelCatalog.swift — Curated registry of downloadable pre-packed Flash-MoE models
|
| 3 |
+
*
|
| 4 |
+
* Each entry maps to a public HuggingFace repo containing pre-packed weights
|
| 5 |
+
* ready for the Flash-MoE engine (config.json, model_weights.bin, packed_experts/).
|
| 6 |
+
*
|
| 7 |
+
* Download URLs: https://huggingface.co/{repoId}/resolve/main/{filename}
|
| 8 |
+
*/
|
| 9 |
+
|
| 10 |
+
import Foundation
|
| 11 |
+
|
| 12 |
+
// MARK: - Data Types
|
| 13 |
+
|
| 14 |
+
struct RepoFile: Identifiable, Codable, Sendable {
|
| 15 |
+
var id: String { filename }
|
| 16 |
+
let filename: String
|
| 17 |
+
let sizeBytes: UInt64
|
| 18 |
+
}
|
| 19 |
+
|
| 20 |
+
struct CatalogEntry: Identifiable, Codable, Sendable {
|
| 21 |
+
let id: String
|
| 22 |
+
let displayName: String
|
| 23 |
+
let repoId: String
|
| 24 |
+
let description: String
|
| 25 |
+
let totalSizeBytes: UInt64
|
| 26 |
+
let quantization: String
|
| 27 |
+
let expertLayers: Int
|
| 28 |
+
let defaultK: Int // Model's native K (experts per token)
|
| 29 |
+
let recommendedK: Int // Recommended K for iPhone (0 = use default)
|
| 30 |
+
let minRAMGB: Int // Minimum RAM in GB
|
| 31 |
+
let files: [RepoFile]
|
| 32 |
+
|
| 33 |
+
var totalSizeGB: Double {
|
| 34 |
+
Double(totalSizeBytes) / (1024.0 * 1024.0 * 1024.0)
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
func downloadURL(for file: RepoFile) -> URL {
|
| 38 |
+
URL(string: "https://huggingface.co/\(repoId)/resolve/main/\(file.filename)")!
|
| 39 |
+
}
|
| 40 |
+
}
|
| 41 |
+
|
| 42 |
+
// MARK: - Curated Model Catalog
|
| 43 |
+
|
| 44 |
+
enum ModelCatalog {
|
| 45 |
+
|
| 46 |
+
/// Pre-packed models available for download.
|
| 47 |
+
/// New models are added in app updates.
|
| 48 |
+
static let models: [CatalogEntry] = [
|
| 49 |
+
// -- Qwen 3.5 35B-A3B 4-bit (full quality) --
|
| 50 |
+
CatalogEntry(
|
| 51 |
+
id: "qwen3.5-35b-a3b-q4",
|
| 52 |
+
displayName: "Qwen 3.5 35B-A3B",
|
| 53 |
+
repoId: "alexintosh/Qwen3.5-35B-A3B-Q4-FlashMoE",
|
| 54 |
+
description: "Compact 35B MoE model. 3B active params per token. Good for 8GB devices.",
|
| 55 |
+
totalSizeBytes: 19_500_000_000,
|
| 56 |
+
quantization: "4-bit",
|
| 57 |
+
expertLayers: 40,
|
| 58 |
+
defaultK: 8,
|
| 59 |
+
recommendedK: 0, // 8 is fine for iPhone
|
| 60 |
+
minRAMGB: 8,
|
| 61 |
+
files: makeFileList(
|
| 62 |
+
configFiles: [
|
| 63 |
+
("config.json", 3_809),
|
| 64 |
+
("model_weights.json", 251_539),
|
| 65 |
+
("model_weights.bin", 1_378_869_376),
|
| 66 |
+
("vocab.bin", 3_360_287),
|
| 67 |
+
("tokenizer.json", 19_989_343),
|
| 68 |
+
("tokenizer.bin", 8_201_040),
|
| 69 |
+
],
|
| 70 |
+
expertLayers: 40,
|
| 71 |
+
expertLayerSize: 452_984_832
|
| 72 |
+
)
|
| 73 |
+
),
|
| 74 |
+
|
| 75 |
+
// -- Qwen 3.5 35B-A3B Tiered (hot=4-bit, cold=2-bit, ~12GB smaller) --
|
| 76 |
+
CatalogEntry(
|
| 77 |
+
id: "qwen3.5-35b-a3b-tiered",
|
| 78 |
+
displayName: "Qwen 3.5 35B-A3B Tiered",
|
| 79 |
+
repoId: "alexintosh/Qwen3.5-35B-A3B-Q4-Tiered-FlashMoE",
|
| 80 |
+
description: "Tiered quantization: hot experts 4-bit, cold 2-bit. ~12GB experts (vs 18GB full). Faster with slight quality trade-off.",
|
| 81 |
+
totalSizeBytes: 13_424_643_082,
|
| 82 |
+
quantization: "tiered (4-bit/2-bit)",
|
| 83 |
+
expertLayers: 40,
|
| 84 |
+
defaultK: 8,
|
| 85 |
+
recommendedK: 0,
|
| 86 |
+
minRAMGB: 8,
|
| 87 |
+
files: makeTieredFileList()
|
| 88 |
+
),
|
| 89 |
+
|
| 90 |
+
// 397B entry removed — repo not published yet. Re-add when
|
| 91 |
+
// alexintosh/Qwen3.5-397B-A17B-Q4-FlashMoE is on HuggingFace.
|
| 92 |
+
]
|
| 93 |
+
|
| 94 |
+
// MARK: - Helpers
|
| 95 |
+
|
| 96 |
+
private static func makeTieredFileList() -> [RepoFile] {
|
| 97 |
+
let configFiles: [(String, UInt64)] = [
|
| 98 |
+
("config.json", 3_809),
|
| 99 |
+
("model_weights.json", 251_539),
|
| 100 |
+
("model_weights.bin", 1_378_869_376),
|
| 101 |
+
("vocab.bin", 3_360_287),
|
| 102 |
+
("tokenizer.json", 19_989_343),
|
| 103 |
+
("tokenizer.bin", 8_201_040),
|
| 104 |
+
]
|
| 105 |
+
// Variable-size tiered expert layers (hot=4-bit, cold=2-bit)
|
| 106 |
+
let layerSizes: [UInt64] = [
|
| 107 |
+
337_379_328, 349_175_808, 342_097_920, 331_087_872, 320_077_824,
|
| 108 |
+
301_989_888, 301_989_888, 289_406_976, 285_474_816, 294_125_568,
|
| 109 |
+
305_922_048, 306_708_480, 297_271_296, 293_339_136, 282_329_088,
|
| 110 |
+
288_620_544, 287_834_112, 292_552_704, 280_756_224, 287_834_112,
|
| 111 |
+
282_329_088, 283_115_520, 301_989_888, 305_135_616, 294_125_568,
|
| 112 |
+
294_125_568, 281_542_656, 292_552_704, 296_484_864, 298_844_160,
|
| 113 |
+
289_406_976, 291_766_272, 301_989_888, 302_776_320, 305_135_616,
|
| 114 |
+
300_417_024, 298_057_728, 304_349_184, 301_989_888, 309_854_208,
|
| 115 |
+
]
|
| 116 |
+
var files = configFiles.map { RepoFile(filename: $0.0, sizeBytes: $0.1) }
|
| 117 |
+
files.append(RepoFile(filename: "packed_experts_tiered/tiered_manifest.json", sizeBytes: 1_005_120))
|
| 118 |
+
for (i, size) in layerSizes.enumerated() {
|
| 119 |
+
files.append(RepoFile(
|
| 120 |
+
filename: String(format: "packed_experts_tiered/layer_%02d.bin", i),
|
| 121 |
+
sizeBytes: size
|
| 122 |
+
))
|
| 123 |
+
}
|
| 124 |
+
return files
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
private static func makeFileList(
|
| 128 |
+
configFiles: [(String, UInt64)],
|
| 129 |
+
expertLayers: Int,
|
| 130 |
+
expertLayerSize: UInt64
|
| 131 |
+
) -> [RepoFile] {
|
| 132 |
+
var files = configFiles.map { RepoFile(filename: $0.0, sizeBytes: $0.1) }
|
| 133 |
+
for i in 0..<expertLayers {
|
| 134 |
+
files.append(RepoFile(
|
| 135 |
+
filename: String(format: "packed_experts/layer_%02d.bin", i),
|
| 136 |
+
sizeBytes: expertLayerSize
|
| 137 |
+
))
|
| 138 |
+
}
|
| 139 |
+
return files
|
| 140 |
+
}
|
| 141 |
+
}
|
FlashMoE-iOS/REVIEW.md
ADDED
|
@@ -0,0 +1,206 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Flash-MoE iOS: Engineering Review & Roadmap
|
| 2 |
+
|
| 3 |
+
Review of the iOS port after achieving 5.5 tok/s on iPhone 17 with Qwen3.5-35B-A3B.
|
| 4 |
+
|
| 5 |
+
## What's Working Well
|
| 6 |
+
|
| 7 |
+
### Architecture Decisions That Paid Off
|
| 8 |
+
|
| 9 |
+
**Unity build was the right call.** The `#include "infer.m"` approach meant zero fork divergence. Every iOS fix (shader loading, tokenizer paths, KV cache cap) was either a small addition to FlashMoEEngine.m or a backward-compatible change to infer.m. The engine runs identical compute paths on both platforms.
|
| 10 |
+
|
| 11 |
+
**Adaptive context via `os_proc_available_memory()`.** The power-of-2 clamping (512 → 8192) based on 25% of available RAM is elegant. Much better than the original hardcoded 2048 cap. This will automatically scale up on future iPhones with more RAM.
|
| 12 |
+
|
| 13 |
+
**`GPU_KV_SEQ` changed from `#define` to `static int`.** This was necessary — the old `#define GPU_KV_SEQ 8192` over-allocated Metal KV buffers regardless of actual context length. The iOS cap to 2048 would still waste 4× the memory on GPU-side KV caches. This fix saves ~150 MB of Metal buffer memory on a typical iPhone load.
|
| 14 |
+
|
| 15 |
+
**KV cache reuse (`flashmoe_generate_continuation`).** This is the single biggest UX feature. Without it, a 5-turn conversation would re-prefill all previous turns every message — at 5.5 tok/s, that's seconds of latency just to re-process history. The `-2` return code for "context full" is a clean signal to the Swift layer.
|
| 16 |
+
|
| 17 |
+
**The profiler overlay.** `os_proc_available_memory()`, `mach_task_basic_info`, `ProcessInfo.thermalState` — all public APIs, no entitlements. This is exactly what you need when debugging memory pressure on a phone running a 35B model.
|
| 18 |
+
|
| 19 |
+
## Concerns
|
| 20 |
+
|
| 21 |
+
### 1. Memory Safety: Silent NULL from calloc
|
| 22 |
+
|
| 23 |
+
**Severity: High**
|
| 24 |
+
|
| 25 |
+
iOS doesn't overcommit memory the way macOS does. When `calloc` fails on iPhone (because you're close to the ~6 GB app limit), it returns NULL silently. The engine has zero NULL checks after calloc/malloc in critical paths:
|
| 26 |
+
|
| 27 |
+
```c
|
| 28 |
+
// FlashMoEEngine.m line 266-267
|
| 29 |
+
ctx->hidden = calloc(cfg.hidden_dim, sizeof(float));
|
| 30 |
+
ctx->logits = calloc(cfg.vocab_size, sizeof(float));
|
| 31 |
+
// If either is NULL → immediate EXC_BAD_ACCESS in the next forward pass
|
| 32 |
+
```
|
| 33 |
+
|
| 34 |
+
Same in `kv_cache_new()`, `linear_attn_state_new()`, `metal_setup()` buffer allocations, and `embed_batch` malloc in the generation loop.
|
| 35 |
+
|
| 36 |
+
**Fix:** Add NULL checks after every allocation in `flashmoe_load()`. If any fail, set `ctx->last_error` and return -1. Also add a pre-flight memory check: compute total expected allocation before starting, compare to `os_proc_available_memory()`, and fail early with a clear error message.
|
| 37 |
+
|
| 38 |
+
### 2. Metal Buffer Allocation Without Error Handling
|
| 39 |
+
|
| 40 |
+
**Severity: High**
|
| 41 |
+
|
| 42 |
+
`metal_setup()` allocates ~40+ Metal buffers. On iPhone, Metal can return `nil` for buffer allocations when GPU memory is exhausted. None of these are checked:
|
| 43 |
+
|
| 44 |
+
```c
|
| 45 |
+
ctx->buf_input = [ctx->device newBufferWithLength:max_in options:MTLResourceStorageModeShared];
|
| 46 |
+
// If nil → next dispatch that uses buf_input crashes
|
| 47 |
+
```
|
| 48 |
+
|
| 49 |
+
The multi-expert buffers are the biggest concern — `MAX_K=8` double-buffered expert data buffers at `expert_alloc_size` each (rounded up to 2 MB alignment). For the 35B model that's `8 × 2 × ~7 MB = ~112 MB` just for expert data buffers.
|
| 50 |
+
|
| 51 |
+
**Fix:** Check every `newBufferWithLength:` return. If any critical buffer fails, log which one, free what was allocated, and return NULL from `metal_setup()`. Consider reducing MAX_K on iPhone (the 35B model uses K=8, but on a memory-constrained device, K=4 with quality degradation might be preferable to a crash).
|
| 52 |
+
|
| 53 |
+
### 3. posix_memalign in metal_setup May Fail on iOS
|
| 54 |
+
|
| 55 |
+
**Severity: Medium**
|
| 56 |
+
|
| 57 |
+
```c
|
| 58 |
+
posix_memalign(&aligned_data, 2*1024*1024, expert_alloc_size);
|
| 59 |
+
```
|
| 60 |
+
|
| 61 |
+
This 2 MB alignment is optimal for DMA on MacBook NVMe. On iPhone:
|
| 62 |
+
- The NVMe controller may have different alignment preferences
|
| 63 |
+
- `posix_memalign` can fail (return non-zero) and leave the pointer uninitialized
|
| 64 |
+
- 2 MB pages may not even be available under memory pressure
|
| 65 |
+
|
| 66 |
+
**Fix:** Check the return value. Consider whether 2 MB alignment matters on iPhone NVMe — it might, or the optimal alignment might be different (64 KB is iOS's large page size). Profile with and without to see if it makes a difference on A18 Pro.
|
| 67 |
+
|
| 68 |
+
### 4. Thread Safety of Static Globals
|
| 69 |
+
|
| 70 |
+
**Severity: Medium**
|
| 71 |
+
|
| 72 |
+
The engine uses ~30 static globals (`cfg`, `g_metal`, `g_deferred`, `g_use_tiered`, etc.). The `FlashMoEContext` struct wraps some state, but the actual compute still reads/writes the globals. This means:
|
| 73 |
+
|
| 74 |
+
- Two `FlashMoEContext` instances would stomp each other's state
|
| 75 |
+
- Calling `flashmoe_generate` from one thread while `flashmoe_get_stats` reads from another could race on `g_deferred`
|
| 76 |
+
|
| 77 |
+
This is fine today (single-instance design), but it's a landmine. The `@unchecked Sendable` on `FlashMoEEngine` in Swift acknowledges this.
|
| 78 |
+
|
| 79 |
+
**Recommendation:** Add a static mutex or `dispatch_once` guard in `flashmoe_create()` that enforces single-instance. If someone calls `flashmoe_create()` while another context exists, return NULL with an error.
|
| 80 |
+
|
| 81 |
+
### 5. Conversation State Leak on Context Full
|
| 82 |
+
|
| 83 |
+
**Severity: Medium**
|
| 84 |
+
|
| 85 |
+
When `flashmoe_generate_continuation` returns -2 (context full), the Swift bridge falls back to `flashmoe_generate` (full re-prefill). But the KV caches and delta-net state still contain the old conversation. The `flashmoe_generate` function resets everything, so this is _probably_ fine, but there's a window where:
|
| 86 |
+
|
| 87 |
+
1. Continuation returns -2
|
| 88 |
+
2. Swift calls `flashmoe_generate` with full chat template
|
| 89 |
+
3. `flashmoe_generate` resets state at line 389-395
|
| 90 |
+
4. New conversation starts clean
|
| 91 |
+
|
| 92 |
+
The concern is if the Swift layer doesn't handle -2 correctly and tries another continuation. Add a guard: after returning -2, set a flag that blocks further continuations until a full generate or reset.
|
| 93 |
+
|
| 94 |
+
### 6. malloc/free per Token in Generation Loop
|
| 95 |
+
|
| 96 |
+
**Severity: Low (performance)**
|
| 97 |
+
|
| 98 |
+
Every token in the generation loop does:
|
| 99 |
+
```c
|
| 100 |
+
float *normed = malloc(cfg.hidden_dim * sizeof(float));
|
| 101 |
+
cpu_rms_norm(ctx->hidden, ctx->final_norm_w, normed, cfg.hidden_dim, cfg.rms_norm_eps);
|
| 102 |
+
memcpy(ctx->hidden, normed, cfg.hidden_dim * sizeof(float));
|
| 103 |
+
free(normed);
|
| 104 |
+
```
|
| 105 |
+
|
| 106 |
+
That's a `malloc(8192)` + `free()` per token. On macOS with the magazine allocator this is ~20 ns, but on iOS under memory pressure, the allocator may need to do more work. At 5.5 tok/s this is ~180 ms per token — a 20 ns malloc is negligible. But it's easy to fix.
|
| 107 |
+
|
| 108 |
+
**Fix:** Pre-allocate `normed` once in `FlashMoEContext` alongside `hidden` and reuse it. Same for `embed_batch` — allocate max size once at load time.
|
| 109 |
+
|
| 110 |
+
### 7. Download Manager: No Integrity Verification
|
| 111 |
+
|
| 112 |
+
**Severity: Medium**
|
| 113 |
+
|
| 114 |
+
Downloaded files are used directly without checksum verification. A corrupted download (network interruption, disk error) would produce garbage inference results or crashes deep in the Metal pipeline. Users would see nonsense output and blame the model.
|
| 115 |
+
|
| 116 |
+
**Fix:** Add SHA-256 checksums to `ModelCatalog.swift` (HuggingFace provides these in the repo metadata). Verify after each file download. Re-download if mismatch.
|
| 117 |
+
|
| 118 |
+
### 8. F_RDAHEAD Disable on iOS
|
| 119 |
+
|
| 120 |
+
**Severity: Low (investigate)**
|
| 121 |
+
|
| 122 |
+
```c
|
| 123 |
+
fcntl(ctx->layer_fds[i], F_RDAHEAD, 0);
|
| 124 |
+
```
|
| 125 |
+
|
| 126 |
+
This was tested on MacBook where expert reads are random and readahead wastes SSD bandwidth. On iPhone, the SSD controller and page cache behavior may differ. The A18 Pro's NVMe controller might benefit from readahead for the sequential-within-expert reads (each expert is ~7 MB of sequential data). Worth A/B testing.
|
| 127 |
+
|
| 128 |
+
## Feature Ideas
|
| 129 |
+
|
| 130 |
+
### Priority 1: Ship Quality
|
| 131 |
+
|
| 132 |
+
1. **Memory pre-flight check.** Before loading, compute total expected allocation (weights mmap + Metal buffers + KV caches + working buffers) and compare to available memory. Show a clear error: "This model needs ~2.8 GB but only 2.1 GB is available. Close other apps or try the Tiered model (saves 34%)."
|
| 133 |
+
|
| 134 |
+
2. **Tiered as default recommendation.** The tiered model is 13.4 GB vs 19.5 GB — that's 6 GB less download and better page cache utilization on iPhone's smaller RAM. The catalog should surface this prominently. Users on iPhone 15 Pro (8 GB) should probably always use tiered.
|
| 135 |
+
|
| 136 |
+
3. **Download resume.** If the app is killed mid-download, resume from where it left off. `URLSession` background downloads support this natively — the `DownloadManager` should persist download progress and use `Range` headers.
|
| 137 |
+
|
| 138 |
+
### Priority 2: Performance
|
| 139 |
+
|
| 140 |
+
4. **Benchmark F_RDAHEAD on iPhone.** The MacBook conclusion ("disable readahead") may not hold on A18. Run 100 tokens with and without `F_RDAHEAD` and compare tok/s.
|
| 141 |
+
|
| 142 |
+
5. **Profile Metal occupancy on A18 GPU.** The half-precision `x_shared` optimization was tuned for M3 Max (40 cores). The A18 Pro has 6 GPU cores with different shared memory and occupancy characteristics. The optimal threadgroup size and shared memory usage may differ. Use Metal System Trace to check.
|
| 143 |
+
|
| 144 |
+
6. **Expert buffer count tuning.** MAX_K=8 allocates 8 double-buffered expert data slots (16 total). The 35B model uses K=8, but if memory is tight, consider K=4 with top-4 routing (lose some quality, save ~56 MB of Metal buffers).
|
| 145 |
+
|
| 146 |
+
### Priority 3: UX Polish
|
| 147 |
+
|
| 148 |
+
7. **Streaming think/reply separation.** The current `<think>` parsing works on completed text but can flicker during streaming. Consider buffering the think block until `</think>` is seen, then revealing it as a disclosure group.
|
| 149 |
+
|
| 150 |
+
8. **Token-level latency display.** The profiler shows average tok/s, but users care about _consistency_. Show a mini sparkline of per-token latency — this would reveal if certain layers have cache misses (spiky) vs warm cache (smooth).
|
| 151 |
+
|
| 152 |
+
9. **Model size warning.** Before downloading a 19.5 GB model, check available disk space and warn if it won't fit. iOS doesn't handle "disk full" gracefully.
|
| 153 |
+
|
| 154 |
+
10. **Background inference continuation.** When the app goes to background during generation, iOS suspends the process. The generation resumes when foregrounded, but the tok/s stat is wrong (counts suspended time). Detect `UIApplication.didEnterBackgroundNotification` and pause the timer.
|
| 155 |
+
|
| 156 |
+
### Priority 4: Future
|
| 157 |
+
|
| 158 |
+
11. **Smaller models.** The 35B-A3B at 19.5 GB is a hard sell for most users. If/when Qwen releases a 7B or 14B MoE variant, that would be the mass-market iOS model. The engine already auto-detects architecture from config.json.
|
| 159 |
+
|
| 160 |
+
12. **On-device expert profiling for tiered.** Run `profile_experts.py` logic on-device after the first few conversations to identify hot experts for _this user's_ workload. Requantize cold experts to 2-bit in the background. Personalized tiered quantization.
|
| 161 |
+
|
| 162 |
+
13. **Widget / Live Activity.** Show generation progress as a Live Activity on the lock screen. "Flash-MoE: generating... 142 tokens, 5.3 tok/s". Great for virality.
|
| 163 |
+
|
| 164 |
+
14. **Shortcuts integration.** Expose `flashmoe_generate` as a Siri Shortcut action. "Ask Flash-MoE: what's the weather like?" — runs inference locally, returns the answer to Shortcuts.
|
| 165 |
+
|
| 166 |
+
## Additional Concerns (from laptop-side review)
|
| 167 |
+
|
| 168 |
+
### 9. No Response to iOS Memory Warnings
|
| 169 |
+
|
| 170 |
+
The app doesn't observe `UIApplication.didReceiveMemoryWarningNotification`. When iOS sends this (which it will — we're using ~2.5GB), the app should at minimum cancel in-flight generation and log the event. Without it, iOS will jettison the app with no recovery. This is the #1 crash risk in the field.
|
| 171 |
+
|
| 172 |
+
**Fix**: In `AppDelegate.swift` or via `NotificationCenter`, observe memory warnings. Cancel generation, optionally show an alert, and consider releasing mmap'd expert data (re-openable on demand).
|
| 173 |
+
|
| 174 |
+
### 10. Thermal Throttling Goes Unhandled
|
| 175 |
+
|
| 176 |
+
The profiler shows thermal state but the engine doesn't react to it. Sustained inference at 5.5 tok/s will push iPhone 15 Pro into `.serious` thermal state within 2-3 minutes. Apple throttles both GPU and CPU at this point, creating a death spiral (slower → longer generation → more heat).
|
| 177 |
+
|
| 178 |
+
**Fix**: When `ProcessInfo.thermalState >= .serious`, insert a small delay between tokens (100-500ms) or pause with a "Phone is cooling down" overlay. This prevents the thermal runaway and gives a better user experience than silently degrading to 2 tok/s.
|
| 179 |
+
|
| 180 |
+
### 11. EOS Token Not Fed Back on Continuation
|
| 181 |
+
|
| 182 |
+
When generation stops on EOS (`<|im_end|>`), the EOS token is detected and the loop breaks — but that token's embedding was never fed through the model's forward pass. The KV cache doesn't have the EOS token's contribution. When `flashmoe_generate_continuation` resumes with `\n<|im_start|>user\n`, the model sees a conversation where the assistant turn didn't properly end.
|
| 183 |
+
|
| 184 |
+
This may or may not matter in practice (the model might be robust to it), but it's technically incorrect. Fix: after the generation loop breaks on EOS, feed the EOS token through one final forward pass (embed + all layers + discard) to properly close the turn in the KV cache.
|
| 185 |
+
|
| 186 |
+
### 12. Share Sheet for Viral Moments
|
| 187 |
+
|
| 188 |
+
Add a "Share" button on each assistant message that creates a formatted card: "Generated locally on iPhone with Flash-MoE • 5.5 tok/s • 35B parameters • No cloud". This is free viral marketing.
|
| 189 |
+
|
| 190 |
+
### 13. Free Space Check Before Download
|
| 191 |
+
|
| 192 |
+
The ModelListView shows model size (19.5 GB) but not device free space. Users start a 19.5GB download without knowing if they have room. Add `FileManager.default.attributesOfFileSystem(forPath:)` to show "X GB free" near the download button.
|
| 193 |
+
|
| 194 |
+
## Performance Notes
|
| 195 |
+
|
| 196 |
+
| Metric | MacBook Pro M3 Max | iPhone 17 (A19) | Ratio |
|
| 197 |
+
|--------|-------------------|-----------------|-------|
|
| 198 |
+
| tok/s | 9.7 | 5.5 | 57% |
|
| 199 |
+
| SSD bandwidth | 17.5 GB/s | ~3-4 GB/s (est.) | ~20% |
|
| 200 |
+
| GPU cores | 40 | 6 | 15% |
|
| 201 |
+
| Memory | 48 GB | 8 GB | 17% |
|
| 202 |
+
| Expert size (4-bit) | 18.1 GB | 18.1 GB | 100% |
|
| 203 |
+
|
| 204 |
+
The iPhone achieves 57% of MacBook speed despite having 15% of the GPU cores and 20% of the SSD bandwidth. This suggests the bottleneck on iPhone is _not_ GPU compute (the 6 cores are fast enough) but SSD streaming. The expert reads (~7 MB × K=8 × 40 layers = 2.24 GB/token) at ~3 GB/s would give ~1.3 tok/s if purely I/O bound — the 5.5 tok/s means the OS page cache is hitting ~75%+ of expert reads, similar to the MacBook's ~71% hit rate. The "Trust the OS" principle holds on iOS.
|
| 205 |
+
|
| 206 |
+
The tiered model should be even faster on iPhone — smaller expert files mean higher page cache hit rate with only 8 GB of RAM managing the cache.
|
FlashMoE-iOS/Resources/Assets.xcassets/AppIcon.appiconset/Contents.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"images" : [
|
| 3 |
+
{
|
| 4 |
+
"idiom" : "universal",
|
| 5 |
+
"platform" : "ios",
|
| 6 |
+
"size" : "1024x1024"
|
| 7 |
+
}
|
| 8 |
+
],
|
| 9 |
+
"info" : {
|
| 10 |
+
"author" : "xcode",
|
| 11 |
+
"version" : 1
|
| 12 |
+
}
|
| 13 |
+
}
|
FlashMoE-iOS/Resources/Assets.xcassets/Contents.json
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"info" : {
|
| 3 |
+
"author" : "xcode",
|
| 4 |
+
"version" : 1
|
| 5 |
+
}
|
| 6 |
+
}
|
FlashMoE-iOS/Resources/FlashMoE.entitlements
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<?xml version="1.0" encoding="UTF-8"?>
|
| 2 |
+
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
|
| 3 |
+
<plist version="1.0">
|
| 4 |
+
<dict>
|
| 5 |
+
<key>com.apple.security.app-sandbox</key>
|
| 6 |
+
<true/>
|
| 7 |
+
<key>com.apple.developer.kernel.extended-virtual-addressing</key>
|
| 8 |
+
<true/>
|
| 9 |
+
<key>com.apple.developer.kernel.increased-memory-limit</key>
|
| 10 |
+
<true/>
|
| 11 |
+
<key>com.apple.security.files.user-selected.read-write</key>
|
| 12 |
+
<true/>
|
| 13 |
+
<key>com.apple.security.network.client</key>
|
| 14 |
+
<true/>
|
| 15 |
+
</dict>
|
| 16 |
+
</plist>
|
FlashMoE-iOS/Services/DownloadManager.swift
ADDED
|
@@ -0,0 +1,512 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* DownloadManager.swift — Background download orchestration for HuggingFace models
|
| 3 |
+
*
|
| 4 |
+
* Uses URLSession background downloads that survive app termination.
|
| 5 |
+
* Downloads files sequentially within a model for clean progress tracking.
|
| 6 |
+
* State persisted to downloads.json for resume across app launches.
|
| 7 |
+
*/
|
| 8 |
+
|
| 9 |
+
import Foundation
|
| 10 |
+
import Observation
|
| 11 |
+
|
| 12 |
+
// MARK: - Download State
|
| 13 |
+
|
| 14 |
+
enum DownloadStatus: String, Codable, Sendable {
|
| 15 |
+
case downloading
|
| 16 |
+
case paused
|
| 17 |
+
case failed
|
| 18 |
+
case complete
|
| 19 |
+
}
|
| 20 |
+
|
| 21 |
+
struct DownloadState: Codable {
|
| 22 |
+
let catalogId: String
|
| 23 |
+
let repoId: String
|
| 24 |
+
var completedFiles: [String]
|
| 25 |
+
var completedBytes: UInt64
|
| 26 |
+
var currentFile: String?
|
| 27 |
+
var status: DownloadStatus
|
| 28 |
+
var errorMessage: String?
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
// MARK: - DownloadManager
|
| 32 |
+
|
| 33 |
+
@Observable
|
| 34 |
+
final class DownloadManager: NSObject, @unchecked Sendable {
|
| 35 |
+
static let shared = DownloadManager()
|
| 36 |
+
|
| 37 |
+
// Observable state
|
| 38 |
+
private(set) var activeDownload: DownloadState?
|
| 39 |
+
private(set) var overallProgress: Double = 0
|
| 40 |
+
private(set) var currentFileProgress: Double = 0
|
| 41 |
+
private(set) var bytesDownloaded: UInt64 = 0
|
| 42 |
+
private(set) var totalBytes: UInt64 = 0
|
| 43 |
+
private(set) var error: String?
|
| 44 |
+
private(set) var downloadSpeed: Double = 0 // bytes/sec
|
| 45 |
+
|
| 46 |
+
// Background session callback
|
| 47 |
+
var backgroundCompletionHandler: (() -> Void)?
|
| 48 |
+
|
| 49 |
+
// Private state
|
| 50 |
+
private var backgroundSession: URLSession!
|
| 51 |
+
private var currentTask: URLSessionDownloadTask?
|
| 52 |
+
private var currentEntry: CatalogEntry?
|
| 53 |
+
private var resumeData: Data?
|
| 54 |
+
private var speedSampleTime: Date?
|
| 55 |
+
private var speedSampleBytes: UInt64 = 0
|
| 56 |
+
|
| 57 |
+
private static let sessionIdentifier = "com.flashmoe.model-download"
|
| 58 |
+
|
| 59 |
+
// MARK: - Initialization
|
| 60 |
+
|
| 61 |
+
override private init() {
|
| 62 |
+
super.init()
|
| 63 |
+
let config = URLSessionConfiguration.background(withIdentifier: Self.sessionIdentifier)
|
| 64 |
+
config.isDiscretionary = false
|
| 65 |
+
config.sessionSendsLaunchEvents = true
|
| 66 |
+
config.allowsCellularAccess = true
|
| 67 |
+
backgroundSession = URLSession(configuration: config, delegate: self, delegateQueue: nil)
|
| 68 |
+
|
| 69 |
+
// Restore persisted state
|
| 70 |
+
loadPersistedState()
|
| 71 |
+
|
| 72 |
+
// Reconnect to any in-flight background tasks
|
| 73 |
+
backgroundSession.getTasksWithCompletionHandler { [weak self] _, _, downloadTasks in
|
| 74 |
+
if let task = downloadTasks.first {
|
| 75 |
+
self?.currentTask = task
|
| 76 |
+
}
|
| 77 |
+
}
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
// MARK: - Public API
|
| 81 |
+
|
| 82 |
+
func startDownload(entry: CatalogEntry) {
|
| 83 |
+
// Allow starting if no active download, or previous one finished/failed
|
| 84 |
+
if let status = activeDownload?.status, status == .downloading || status == .paused {
|
| 85 |
+
if activeDownload?.catalogId != entry.id {
|
| 86 |
+
error = "A different download is already in progress"
|
| 87 |
+
return
|
| 88 |
+
}
|
| 89 |
+
// Same model — resume instead
|
| 90 |
+
resumeDownload()
|
| 91 |
+
return
|
| 92 |
+
}
|
| 93 |
+
|
| 94 |
+
// Clear stale state from previous download
|
| 95 |
+
if activeDownload != nil {
|
| 96 |
+
activeDownload = nil
|
| 97 |
+
clearPersistedState()
|
| 98 |
+
}
|
| 99 |
+
|
| 100 |
+
// Check disk space
|
| 101 |
+
let available = availableDiskSpace()
|
| 102 |
+
if available < entry.totalSizeBytes {
|
| 103 |
+
let needed = formatBytes(entry.totalSizeBytes)
|
| 104 |
+
let have = formatBytes(available)
|
| 105 |
+
error = "Not enough space: \(needed) needed, \(have) available"
|
| 106 |
+
return
|
| 107 |
+
}
|
| 108 |
+
|
| 109 |
+
error = nil
|
| 110 |
+
currentEntry = entry
|
| 111 |
+
totalBytes = entry.totalSizeBytes
|
| 112 |
+
|
| 113 |
+
// Create model directory
|
| 114 |
+
let modelDir = modelDirectory(for: entry.id)
|
| 115 |
+
createDirectoryStructure(for: entry, at: modelDir)
|
| 116 |
+
|
| 117 |
+
activeDownload = DownloadState(
|
| 118 |
+
catalogId: entry.id,
|
| 119 |
+
repoId: entry.repoId,
|
| 120 |
+
completedFiles: [],
|
| 121 |
+
completedBytes: 0,
|
| 122 |
+
currentFile: nil,
|
| 123 |
+
status: .downloading
|
| 124 |
+
)
|
| 125 |
+
persistState()
|
| 126 |
+
downloadNextFile()
|
| 127 |
+
}
|
| 128 |
+
|
| 129 |
+
func pauseDownload() {
|
| 130 |
+
guard activeDownload?.status == .downloading else { return }
|
| 131 |
+
|
| 132 |
+
currentTask?.cancel(byProducingResumeData: { [weak self] data in
|
| 133 |
+
guard let self else { return }
|
| 134 |
+
self.resumeData = data
|
| 135 |
+
self.activeDownload?.status = .paused
|
| 136 |
+
self.persistState()
|
| 137 |
+
self.currentTask = nil
|
| 138 |
+
})
|
| 139 |
+
}
|
| 140 |
+
|
| 141 |
+
func resumeDownload() {
|
| 142 |
+
guard activeDownload?.status == .paused || activeDownload?.status == .failed else { return }
|
| 143 |
+
|
| 144 |
+
// Resolve the catalog entry
|
| 145 |
+
if currentEntry == nil, let catalogId = activeDownload?.catalogId {
|
| 146 |
+
currentEntry = ModelCatalog.models.first { $0.id == catalogId }
|
| 147 |
+
}
|
| 148 |
+
guard currentEntry != nil else {
|
| 149 |
+
error = "Cannot find model in catalog"
|
| 150 |
+
return
|
| 151 |
+
}
|
| 152 |
+
|
| 153 |
+
error = nil
|
| 154 |
+
activeDownload?.status = .downloading
|
| 155 |
+
activeDownload?.errorMessage = nil
|
| 156 |
+
totalBytes = currentEntry?.totalSizeBytes ?? 0
|
| 157 |
+
persistState()
|
| 158 |
+
|
| 159 |
+
if let resumeData {
|
| 160 |
+
let task = backgroundSession.downloadTask(withResumeData: resumeData)
|
| 161 |
+
task.resume()
|
| 162 |
+
currentTask = task
|
| 163 |
+
self.resumeData = nil
|
| 164 |
+
} else {
|
| 165 |
+
downloadNextFile()
|
| 166 |
+
}
|
| 167 |
+
}
|
| 168 |
+
|
| 169 |
+
func cancelDownload() {
|
| 170 |
+
currentTask?.cancel()
|
| 171 |
+
currentTask = nil
|
| 172 |
+
resumeData = nil
|
| 173 |
+
|
| 174 |
+
if let catalogId = activeDownload?.catalogId {
|
| 175 |
+
let dir = modelDirectory(for: catalogId)
|
| 176 |
+
try? FileManager.default.removeItem(at: dir)
|
| 177 |
+
}
|
| 178 |
+
|
| 179 |
+
activeDownload = nil
|
| 180 |
+
overallProgress = 0
|
| 181 |
+
currentFileProgress = 0
|
| 182 |
+
bytesDownloaded = 0
|
| 183 |
+
totalBytes = 0
|
| 184 |
+
error = nil
|
| 185 |
+
currentEntry = nil
|
| 186 |
+
clearPersistedState()
|
| 187 |
+
}
|
| 188 |
+
|
| 189 |
+
func deleteModel(catalogId: String) {
|
| 190 |
+
let dir = modelDirectory(for: catalogId)
|
| 191 |
+
try? FileManager.default.removeItem(at: dir)
|
| 192 |
+
|
| 193 |
+
if activeDownload?.catalogId == catalogId {
|
| 194 |
+
activeDownload = nil
|
| 195 |
+
clearPersistedState()
|
| 196 |
+
}
|
| 197 |
+
}
|
| 198 |
+
|
| 199 |
+
func isModelDownloaded(_ catalogId: String) -> Bool {
|
| 200 |
+
let dir = modelDirectory(for: catalogId)
|
| 201 |
+
return FlashMoEEngine.validateModel(at: dir.path)
|
| 202 |
+
}
|
| 203 |
+
|
| 204 |
+
func modelPath(for catalogId: String) -> String {
|
| 205 |
+
modelDirectory(for: catalogId).path
|
| 206 |
+
}
|
| 207 |
+
|
| 208 |
+
// MARK: - File Management
|
| 209 |
+
|
| 210 |
+
private func modelDirectory(for catalogId: String) -> URL {
|
| 211 |
+
let docs = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)[0]
|
| 212 |
+
return docs.appendingPathComponent(catalogId)
|
| 213 |
+
}
|
| 214 |
+
|
| 215 |
+
private func createDirectoryStructure(for entry: CatalogEntry, at baseURL: URL) {
|
| 216 |
+
let fm = FileManager.default
|
| 217 |
+
try? fm.createDirectory(at: baseURL, withIntermediateDirectories: true)
|
| 218 |
+
|
| 219 |
+
// Create subdirectories for expert files
|
| 220 |
+
var subdirs = Set<String>()
|
| 221 |
+
for file in entry.files {
|
| 222 |
+
let url = baseURL.appendingPathComponent(file.filename)
|
| 223 |
+
let parent = url.deletingLastPathComponent()
|
| 224 |
+
if parent != baseURL {
|
| 225 |
+
subdirs.insert(parent.path)
|
| 226 |
+
}
|
| 227 |
+
}
|
| 228 |
+
for dir in subdirs {
|
| 229 |
+
try? fm.createDirectory(atPath: dir, withIntermediateDirectories: true)
|
| 230 |
+
}
|
| 231 |
+
}
|
| 232 |
+
|
| 233 |
+
// MARK: - Sequential Download Engine
|
| 234 |
+
|
| 235 |
+
private func downloadNextFile() {
|
| 236 |
+
guard var state = activeDownload, let entry = currentEntry else { return }
|
| 237 |
+
|
| 238 |
+
// Find next file to download
|
| 239 |
+
let nextFile = entry.files.first { !state.completedFiles.contains($0.filename) }
|
| 240 |
+
|
| 241 |
+
guard let file = nextFile else {
|
| 242 |
+
// All files downloaded
|
| 243 |
+
state.status = .complete
|
| 244 |
+
state.currentFile = nil
|
| 245 |
+
activeDownload = state
|
| 246 |
+
overallProgress = 1.0
|
| 247 |
+
persistState()
|
| 248 |
+
|
| 249 |
+
// Protect from iOS storage optimization (exclude from backup/purge)
|
| 250 |
+
let dir = modelDirectory(for: entry.id)
|
| 251 |
+
var dirURL = dir
|
| 252 |
+
var values = URLResourceValues()
|
| 253 |
+
values.isExcludedFromBackup = true
|
| 254 |
+
try? dirURL.setResourceValues(values)
|
| 255 |
+
if let enumerator = FileManager.default.enumerator(at: dir, includingPropertiesForKeys: nil) {
|
| 256 |
+
while let fileURL = enumerator.nextObject() as? URL {
|
| 257 |
+
var fURL = fileURL
|
| 258 |
+
try? fURL.setResourceValues(values)
|
| 259 |
+
}
|
| 260 |
+
}
|
| 261 |
+
|
| 262 |
+
// Validate the model
|
| 263 |
+
if !FlashMoEEngine.validateModel(at: dir.path) {
|
| 264 |
+
error = "Download complete but model validation failed"
|
| 265 |
+
state.status = .failed
|
| 266 |
+
state.errorMessage = "Validation failed — some files may be corrupt"
|
| 267 |
+
activeDownload = state
|
| 268 |
+
persistState()
|
| 269 |
+
}
|
| 270 |
+
return
|
| 271 |
+
}
|
| 272 |
+
|
| 273 |
+
state.currentFile = file.filename
|
| 274 |
+
activeDownload = state
|
| 275 |
+
persistState()
|
| 276 |
+
|
| 277 |
+
let url = entry.downloadURL(for: file)
|
| 278 |
+
let task = backgroundSession.downloadTask(with: url)
|
| 279 |
+
task.taskDescription = file.filename
|
| 280 |
+
task.resume()
|
| 281 |
+
currentTask = task
|
| 282 |
+
currentFileProgress = 0
|
| 283 |
+
speedSampleTime = Date()
|
| 284 |
+
speedSampleBytes = bytesDownloaded
|
| 285 |
+
}
|
| 286 |
+
|
| 287 |
+
// MARK: - Disk Space
|
| 288 |
+
|
| 289 |
+
private func availableDiskSpace() -> UInt64 {
|
| 290 |
+
let docs = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)[0]
|
| 291 |
+
guard let values = try? docs.resourceValues(forKeys: [.volumeAvailableCapacityForImportantUsageKey]),
|
| 292 |
+
let capacity = values.volumeAvailableCapacityForImportantUsage else {
|
| 293 |
+
return 0
|
| 294 |
+
}
|
| 295 |
+
return UInt64(capacity)
|
| 296 |
+
}
|
| 297 |
+
|
| 298 |
+
// MARK: - State Persistence
|
| 299 |
+
|
| 300 |
+
private var stateFileURL: URL {
|
| 301 |
+
let docs = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)[0]
|
| 302 |
+
return docs.appendingPathComponent("downloads.json")
|
| 303 |
+
}
|
| 304 |
+
|
| 305 |
+
private func persistState() {
|
| 306 |
+
guard let state = activeDownload else { return }
|
| 307 |
+
if let data = try? JSONEncoder().encode(state) {
|
| 308 |
+
try? data.write(to: stateFileURL)
|
| 309 |
+
}
|
| 310 |
+
}
|
| 311 |
+
|
| 312 |
+
private func clearPersistedState() {
|
| 313 |
+
try? FileManager.default.removeItem(at: stateFileURL)
|
| 314 |
+
}
|
| 315 |
+
|
| 316 |
+
private func loadPersistedState() {
|
| 317 |
+
guard let data = try? Data(contentsOf: stateFileURL),
|
| 318 |
+
let state = try? JSONDecoder().decode(DownloadState.self, from: data) else {
|
| 319 |
+
return
|
| 320 |
+
}
|
| 321 |
+
|
| 322 |
+
activeDownload = state
|
| 323 |
+
currentEntry = ModelCatalog.models.first { $0.id == state.catalogId }
|
| 324 |
+
|
| 325 |
+
if let entry = currentEntry {
|
| 326 |
+
totalBytes = entry.totalSizeBytes
|
| 327 |
+
bytesDownloaded = state.completedBytes
|
| 328 |
+
overallProgress = totalBytes > 0 ? Double(bytesDownloaded) / Double(totalBytes) : 0
|
| 329 |
+
}
|
| 330 |
+
}
|
| 331 |
+
|
| 332 |
+
// MARK: - Formatting
|
| 333 |
+
|
| 334 |
+
private func formatBytes(_ bytes: UInt64) -> String {
|
| 335 |
+
let gb = Double(bytes) / (1024 * 1024 * 1024)
|
| 336 |
+
if gb >= 1 {
|
| 337 |
+
return String(format: "%.1f GB", gb)
|
| 338 |
+
}
|
| 339 |
+
let mb = Double(bytes) / (1024 * 1024)
|
| 340 |
+
return String(format: "%.0f MB", mb)
|
| 341 |
+
}
|
| 342 |
+
}
|
| 343 |
+
|
| 344 |
+
// MARK: - URLSessionDownloadDelegate
|
| 345 |
+
|
| 346 |
+
extension DownloadManager: URLSessionDownloadDelegate {
|
| 347 |
+
|
| 348 |
+
func urlSession(
|
| 349 |
+
_ session: URLSession,
|
| 350 |
+
downloadTask: URLSessionDownloadTask,
|
| 351 |
+
didFinishDownloadingTo location: URL
|
| 352 |
+
) {
|
| 353 |
+
guard let entry = currentEntry,
|
| 354 |
+
let filename = downloadTask.taskDescription else { return }
|
| 355 |
+
|
| 356 |
+
// Check HTTP status code — HuggingFace returns 200 HTML pages for 404s
|
| 357 |
+
if let httpResponse = downloadTask.response as? HTTPURLResponse,
|
| 358 |
+
httpResponse.statusCode != 200 {
|
| 359 |
+
let statusCode = httpResponse.statusCode
|
| 360 |
+
DispatchQueue.main.async { [weak self] in
|
| 361 |
+
guard let self, var state = self.activeDownload else { return }
|
| 362 |
+
self.error = "HTTP \(statusCode) downloading \(filename)"
|
| 363 |
+
state.status = .failed
|
| 364 |
+
state.errorMessage = self.error
|
| 365 |
+
self.activeDownload = state
|
| 366 |
+
self.persistState()
|
| 367 |
+
}
|
| 368 |
+
return
|
| 369 |
+
}
|
| 370 |
+
|
| 371 |
+
// Move from temp to model directory (must happen synchronously before this method returns)
|
| 372 |
+
let dest = modelDirectory(for: entry.id).appendingPathComponent(filename)
|
| 373 |
+
let fm = FileManager.default
|
| 374 |
+
try? fm.removeItem(at: dest)
|
| 375 |
+
|
| 376 |
+
var moveError: Error?
|
| 377 |
+
do {
|
| 378 |
+
try fm.moveItem(at: location, to: dest)
|
| 379 |
+
} catch {
|
| 380 |
+
moveError = error
|
| 381 |
+
}
|
| 382 |
+
|
| 383 |
+
// Compute file size on this thread before dispatching
|
| 384 |
+
let actualSize: UInt64
|
| 385 |
+
if moveError == nil {
|
| 386 |
+
let attrs = try? fm.attributesOfItem(atPath: dest.path)
|
| 387 |
+
actualSize = attrs?[.size] as? UInt64 ?? 0
|
| 388 |
+
} else {
|
| 389 |
+
actualSize = 0
|
| 390 |
+
}
|
| 391 |
+
|
| 392 |
+
let expectedSize = entry.files.first(where: { $0.filename == filename })?.sizeBytes ?? 0
|
| 393 |
+
|
| 394 |
+
// Check if we got an HTML error page instead of the actual file
|
| 395 |
+
// (HuggingFace sometimes returns 200 with HTML for missing LFS files)
|
| 396 |
+
if actualSize > 0 && actualSize < 10_000 && expectedSize > 100_000 {
|
| 397 |
+
// Downloaded file is suspiciously small — likely an error page
|
| 398 |
+
try? fm.removeItem(at: dest)
|
| 399 |
+
DispatchQueue.main.async { [weak self] in
|
| 400 |
+
guard let self, var state = self.activeDownload else { return }
|
| 401 |
+
self.error = "File \(filename) not found on server (got \(actualSize) bytes, expected \(self.formatBytes(expectedSize)))"
|
| 402 |
+
state.status = .failed
|
| 403 |
+
state.errorMessage = self.error
|
| 404 |
+
self.activeDownload = state
|
| 405 |
+
self.persistState()
|
| 406 |
+
}
|
| 407 |
+
return
|
| 408 |
+
}
|
| 409 |
+
|
| 410 |
+
// All @Observable mutations on main thread
|
| 411 |
+
DispatchQueue.main.async { [weak self] in
|
| 412 |
+
guard let self, var state = self.activeDownload else { return }
|
| 413 |
+
|
| 414 |
+
if let moveError {
|
| 415 |
+
self.error = "Failed to save \(filename): \(moveError.localizedDescription)"
|
| 416 |
+
state.status = .failed
|
| 417 |
+
state.errorMessage = self.error
|
| 418 |
+
self.activeDownload = state
|
| 419 |
+
self.persistState()
|
| 420 |
+
return
|
| 421 |
+
}
|
| 422 |
+
|
| 423 |
+
// Validate file size
|
| 424 |
+
if actualSize > 0 && expectedSize > 0 && actualSize < expectedSize * 9 / 10 {
|
| 425 |
+
self.error = "File \(filename) is too small (\(actualSize) vs expected \(expectedSize))"
|
| 426 |
+
state.status = .failed
|
| 427 |
+
state.errorMessage = self.error
|
| 428 |
+
self.activeDownload = state
|
| 429 |
+
self.persistState()
|
| 430 |
+
return
|
| 431 |
+
}
|
| 432 |
+
|
| 433 |
+
state.completedBytes += actualSize > 0 ? actualSize : expectedSize
|
| 434 |
+
state.completedFiles.append(filename)
|
| 435 |
+
state.currentFile = nil
|
| 436 |
+
self.activeDownload = state
|
| 437 |
+
self.bytesDownloaded = state.completedBytes
|
| 438 |
+
self.persistState()
|
| 439 |
+
|
| 440 |
+
// Start next file
|
| 441 |
+
self.downloadNextFile()
|
| 442 |
+
}
|
| 443 |
+
}
|
| 444 |
+
|
| 445 |
+
func urlSession(
|
| 446 |
+
_ session: URLSession,
|
| 447 |
+
downloadTask: URLSessionDownloadTask,
|
| 448 |
+
didWriteData bytesWritten: Int64,
|
| 449 |
+
totalBytesWritten: Int64,
|
| 450 |
+
totalBytesExpectedToWrite: Int64
|
| 451 |
+
) {
|
| 452 |
+
// Compute values on background thread
|
| 453 |
+
let fileProgress = totalBytesExpectedToWrite > 0
|
| 454 |
+
? Double(totalBytesWritten) / Double(totalBytesExpectedToWrite) : 0
|
| 455 |
+
let completed = activeDownload?.completedBytes ?? 0
|
| 456 |
+
let currentTotal = completed + UInt64(totalBytesWritten)
|
| 457 |
+
let total = totalBytes
|
| 458 |
+
let overall = total > 0 ? Double(currentTotal) / Double(total) : 0
|
| 459 |
+
|
| 460 |
+
// Speed sampling (non-observable state, safe on background)
|
| 461 |
+
var newSpeed: Double?
|
| 462 |
+
if let sampleTime = speedSampleTime, Date().timeIntervalSince(sampleTime) >= 2 {
|
| 463 |
+
let elapsed = Date().timeIntervalSince(sampleTime)
|
| 464 |
+
let delta = currentTotal - speedSampleBytes
|
| 465 |
+
newSpeed = Double(delta) / elapsed
|
| 466 |
+
speedSampleTime = Date()
|
| 467 |
+
speedSampleBytes = currentTotal
|
| 468 |
+
}
|
| 469 |
+
|
| 470 |
+
// All @Observable mutations on main thread
|
| 471 |
+
DispatchQueue.main.async { [weak self] in
|
| 472 |
+
guard let self else { return }
|
| 473 |
+
self.currentFileProgress = fileProgress
|
| 474 |
+
self.bytesDownloaded = currentTotal
|
| 475 |
+
self.overallProgress = overall
|
| 476 |
+
if let newSpeed {
|
| 477 |
+
self.downloadSpeed = newSpeed
|
| 478 |
+
}
|
| 479 |
+
}
|
| 480 |
+
}
|
| 481 |
+
|
| 482 |
+
func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: (any Error)?) {
|
| 483 |
+
guard let error else { return }
|
| 484 |
+
|
| 485 |
+
let nsError = error as NSError
|
| 486 |
+
if nsError.code == NSURLErrorCancelled {
|
| 487 |
+
return
|
| 488 |
+
}
|
| 489 |
+
|
| 490 |
+
// Save resume data if available (non-observable)
|
| 491 |
+
if let data = nsError.userInfo[NSURLSessionDownloadTaskResumeData] as? Data {
|
| 492 |
+
self.resumeData = data
|
| 493 |
+
}
|
| 494 |
+
|
| 495 |
+
let errorMsg = error.localizedDescription
|
| 496 |
+
|
| 497 |
+
DispatchQueue.main.async { [weak self] in
|
| 498 |
+
guard let self else { return }
|
| 499 |
+
self.error = errorMsg
|
| 500 |
+
self.activeDownload?.status = .failed
|
| 501 |
+
self.activeDownload?.errorMessage = errorMsg
|
| 502 |
+
self.persistState()
|
| 503 |
+
}
|
| 504 |
+
}
|
| 505 |
+
|
| 506 |
+
func urlSessionDidFinishEvents(forBackgroundURLSession session: URLSession) {
|
| 507 |
+
DispatchQueue.main.async { [weak self] in
|
| 508 |
+
self?.backgroundCompletionHandler?()
|
| 509 |
+
self?.backgroundCompletionHandler = nil
|
| 510 |
+
}
|
| 511 |
+
}
|
| 512 |
+
}
|
FlashMoE-iOS/Views/BenchmarkView.swift
ADDED
|
@@ -0,0 +1,547 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* BenchmarkView.swift -- Built-in benchmark mode for Flash-MoE iOS
|
| 3 |
+
*
|
| 4 |
+
* Runs a test matrix of prompts x settings configurations,
|
| 5 |
+
* logs structured results (tab-separated), and provides
|
| 6 |
+
* copy-to-clipboard for analysis.
|
| 7 |
+
*/
|
| 8 |
+
|
| 9 |
+
import SwiftUI
|
| 10 |
+
|
| 11 |
+
// MARK: - Benchmark Data Types
|
| 12 |
+
|
| 13 |
+
struct BenchmarkPrompt: Identifiable {
|
| 14 |
+
let id: String
|
| 15 |
+
let label: String
|
| 16 |
+
let text: String
|
| 17 |
+
|
| 18 |
+
static let builtIn: [BenchmarkPrompt] = [
|
| 19 |
+
BenchmarkPrompt(id: "short", label: "Short", text: "Hi"),
|
| 20 |
+
BenchmarkPrompt(id: "medium", label: "Medium", text: "What is an SSD and how does it work?"),
|
| 21 |
+
BenchmarkPrompt(id: "long", label: "Long", text: "Explain the differences between CPU, GPU, and TPU architectures. Compare their memory hierarchies, parallelism models, and ideal workloads. Include specific examples."),
|
| 22 |
+
BenchmarkPrompt(id: "json", label: "JSON", text: "Respond with a JSON object containing: name, age, hobbies (array of 3), and address (nested object with street, city, country)"),
|
| 23 |
+
BenchmarkPrompt(id: "code", label: "Code", text: "Write a Python function that implements binary search on a sorted array. Include docstring and type hints."),
|
| 24 |
+
]
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
struct BenchmarkConfig: Identifiable {
|
| 28 |
+
let id: String
|
| 29 |
+
let name: String
|
| 30 |
+
let activeExpertsK: Int
|
| 31 |
+
let cmdMerge: Bool
|
| 32 |
+
let fusedAttention: Bool
|
| 33 |
+
let cacheIOSplit: Int
|
| 34 |
+
let fp16Accumulation: Bool
|
| 35 |
+
|
| 36 |
+
static let builtIn: [BenchmarkConfig] = [
|
| 37 |
+
BenchmarkConfig(id: "baseline", name: "Baseline K=8", activeExpertsK: 8, cmdMerge: false, fusedAttention: false, cacheIOSplit: 1, fp16Accumulation: false),
|
| 38 |
+
BenchmarkConfig(id: "k4_merge", name: "K=4+Merge", activeExpertsK: 4, cmdMerge: true, fusedAttention: false, cacheIOSplit: 1, fp16Accumulation: false),
|
| 39 |
+
BenchmarkConfig(id: "k4_merge_fused", name: "K=4+Merge+Fused F2", activeExpertsK: 4, cmdMerge: true, fusedAttention: true, cacheIOSplit: 2, fp16Accumulation: false),
|
| 40 |
+
BenchmarkConfig(id: "k4_merge_fused_fp16", name: "K=4+Merge+Fused+FP16", activeExpertsK: 4, cmdMerge: true, fusedAttention: true, cacheIOSplit: 2, fp16Accumulation: true),
|
| 41 |
+
BenchmarkConfig(id: "k3_merge_fused_fp16", name: "K=3+Merge+Fused+FP16", activeExpertsK: 3, cmdMerge: true, fusedAttention: true, cacheIOSplit: 2, fp16Accumulation: true),
|
| 42 |
+
]
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
struct BenchmarkResult: Identifiable {
|
| 46 |
+
let id = UUID()
|
| 47 |
+
let configName: String
|
| 48 |
+
let promptLabel: String
|
| 49 |
+
let promptText: String
|
| 50 |
+
let prefillTokens: Int
|
| 51 |
+
let ttftMs: Double
|
| 52 |
+
let tokensGenerated: Int
|
| 53 |
+
let decodeTokPerSec: Double // actual decode speed (tokens after first / decode time)
|
| 54 |
+
let engineTokPerSec: Double // engine-reported tok/s
|
| 55 |
+
let totalMs: Double
|
| 56 |
+
let outputSnippet: String
|
| 57 |
+
let quality: String
|
| 58 |
+
let thermalState: String
|
| 59 |
+
|
| 60 |
+
var tsvLine: String {
|
| 61 |
+
let truncatedPrompt = promptText.count > 50
|
| 62 |
+
? String(promptText.prefix(50)) + "..."
|
| 63 |
+
: promptText
|
| 64 |
+
let cleanSnippet = outputSnippet
|
| 65 |
+
.replacingOccurrences(of: "\t", with: " ")
|
| 66 |
+
.replacingOccurrences(of: "\n", with: " ")
|
| 67 |
+
return "\(configName)\t\(truncatedPrompt)\t\(prefillTokens)\t\(String(format: "%.0f", ttftMs))\t\(tokensGenerated)\t\(String(format: "%.1f", decodeTokPerSec))\t\(String(format: "%.1f", engineTokPerSec))\t\(String(format: "%.0f", totalMs))\t\(thermalState)\t\(quality)\t\(cleanSnippet)"
|
| 68 |
+
}
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
// MARK: - Benchmark Runner (Observable)
|
| 72 |
+
|
| 73 |
+
@Observable @MainActor
|
| 74 |
+
final class BenchmarkRunner {
|
| 75 |
+
var isRunning = false
|
| 76 |
+
var isCancelled = false
|
| 77 |
+
var currentConfigIndex = 0
|
| 78 |
+
var currentPromptIndex = 0
|
| 79 |
+
var totalConfigs: Int = BenchmarkConfig.builtIn.count
|
| 80 |
+
var totalPrompts: Int = BenchmarkPrompt.builtIn.count
|
| 81 |
+
var results: [BenchmarkResult] = []
|
| 82 |
+
var statusMessage = "Ready"
|
| 83 |
+
var progressDetail = ""
|
| 84 |
+
|
| 85 |
+
let maxTokensPerRun = 50
|
| 86 |
+
let cooldownBetweenPrompts: UInt64 = 5_000_000_000 // 5 seconds
|
| 87 |
+
let cooldownBetweenConfigs: UInt64 = 10_000_000_000 // 10 seconds
|
| 88 |
+
|
| 89 |
+
func cancel() {
|
| 90 |
+
isCancelled = true
|
| 91 |
+
statusMessage = "Cancelling..."
|
| 92 |
+
}
|
| 93 |
+
|
| 94 |
+
private func thermalStateString() -> String {
|
| 95 |
+
switch ProcessInfo.processInfo.thermalState {
|
| 96 |
+
case .nominal: return "nominal"
|
| 97 |
+
case .fair: return "fair"
|
| 98 |
+
case .serious: return "serious"
|
| 99 |
+
case .critical: return "critical"
|
| 100 |
+
@unknown default: return "unknown"
|
| 101 |
+
}
|
| 102 |
+
}
|
| 103 |
+
|
| 104 |
+
private func waitForCooldown(duration: UInt64, label: String) async {
|
| 105 |
+
let thermal = thermalStateString()
|
| 106 |
+
if thermal == "serious" || thermal == "critical" {
|
| 107 |
+
progressDetail = "Thermal \(thermal) — cooling down (30s)..."
|
| 108 |
+
try? await Task.sleep(nanoseconds: 30_000_000_000)
|
| 109 |
+
} else {
|
| 110 |
+
progressDetail = "\(label) cooldown (\(duration / 1_000_000_000)s)..."
|
| 111 |
+
try? await Task.sleep(nanoseconds: duration)
|
| 112 |
+
}
|
| 113 |
+
}
|
| 114 |
+
|
| 115 |
+
func run(engine: FlashMoEEngine, modelPath: String,
|
| 116 |
+
configIds: Set<String>? = nil, promptIds: Set<String>? = nil) async {
|
| 117 |
+
isRunning = true
|
| 118 |
+
isCancelled = false
|
| 119 |
+
results = []
|
| 120 |
+
currentConfigIndex = 0
|
| 121 |
+
currentPromptIndex = 0
|
| 122 |
+
|
| 123 |
+
let configs = configIds != nil
|
| 124 |
+
? BenchmarkConfig.builtIn.filter { configIds!.contains($0.id) }
|
| 125 |
+
: BenchmarkConfig.builtIn
|
| 126 |
+
let prompts = promptIds != nil
|
| 127 |
+
? BenchmarkPrompt.builtIn.filter { promptIds!.contains($0.id) }
|
| 128 |
+
: BenchmarkPrompt.builtIn
|
| 129 |
+
|
| 130 |
+
totalConfigs = configs.count
|
| 131 |
+
totalPrompts = prompts.count
|
| 132 |
+
statusMessage = "Starting benchmark (\(configs.count)×\(prompts.count))..."
|
| 133 |
+
|
| 134 |
+
for (ci, config) in configs.enumerated() {
|
| 135 |
+
if isCancelled { break }
|
| 136 |
+
|
| 137 |
+
currentConfigIndex = ci
|
| 138 |
+
statusMessage = "Config \(ci + 1)/\(configs.count): \(config.name)"
|
| 139 |
+
progressDetail = "Applying settings..."
|
| 140 |
+
|
| 141 |
+
// Apply config by setting C globals directly — NO model reload needed.
|
| 142 |
+
// This keeps the fullScreenCover alive and avoids 5-10s reload per config.
|
| 143 |
+
engine.applyBenchmarkConfig(
|
| 144 |
+
activeExpertsK: config.activeExpertsK,
|
| 145 |
+
cmdMerge: config.cmdMerge,
|
| 146 |
+
fusedAttention: config.fusedAttention,
|
| 147 |
+
cacheIOSplit: config.cacheIOSplit,
|
| 148 |
+
fp16Accumulation: config.fp16Accumulation
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
// Run each prompt with this config
|
| 152 |
+
for (pi, prompt) in prompts.enumerated() {
|
| 153 |
+
if isCancelled { break }
|
| 154 |
+
|
| 155 |
+
currentPromptIndex = pi
|
| 156 |
+
statusMessage = "Config \(ci + 1)/\(configs.count), Prompt \(pi + 1)/\(prompts.count)"
|
| 157 |
+
progressDetail = "\(config.name) | \(prompt.label)"
|
| 158 |
+
|
| 159 |
+
// Reset conversation state between prompts
|
| 160 |
+
engine.reset()
|
| 161 |
+
|
| 162 |
+
let result = await runSingleBenchmark(
|
| 163 |
+
engine: engine,
|
| 164 |
+
config: config,
|
| 165 |
+
prompt: prompt
|
| 166 |
+
)
|
| 167 |
+
|
| 168 |
+
results.append(result)
|
| 169 |
+
|
| 170 |
+
// Cooldown between prompts
|
| 171 |
+
if pi < prompts.count - 1 && !isCancelled {
|
| 172 |
+
await waitForCooldown(duration: cooldownBetweenPrompts, label: "Prompt")
|
| 173 |
+
}
|
| 174 |
+
}
|
| 175 |
+
|
| 176 |
+
// Cooldown between configs
|
| 177 |
+
if ci < configs.count - 1 && !isCancelled {
|
| 178 |
+
await waitForCooldown(duration: cooldownBetweenConfigs, label: "Config")
|
| 179 |
+
}
|
| 180 |
+
}
|
| 181 |
+
|
| 182 |
+
isRunning = false
|
| 183 |
+
statusMessage = isCancelled ? "Cancelled" : "Complete (\(results.count) runs)"
|
| 184 |
+
progressDetail = ""
|
| 185 |
+
}
|
| 186 |
+
|
| 187 |
+
private func runSingleBenchmark(
|
| 188 |
+
engine: FlashMoEEngine,
|
| 189 |
+
config: BenchmarkConfig,
|
| 190 |
+
prompt: BenchmarkPrompt
|
| 191 |
+
) async -> BenchmarkResult {
|
| 192 |
+
// Build Qwen chat template (no thinking)
|
| 193 |
+
let formatted = "<|im_start|>system\nYou are a helpful assistant. /no_think<|im_end|>\n<|im_start|>user\n\(prompt.text)<|im_end|>\n<|im_start|>assistant\n<think>\n"
|
| 194 |
+
|
| 195 |
+
let thermal = thermalStateString()
|
| 196 |
+
let startTime = CFAbsoluteTimeGetCurrent()
|
| 197 |
+
var firstTokenTime: Double = 0
|
| 198 |
+
var output = ""
|
| 199 |
+
var tokenCount = 0
|
| 200 |
+
var lastEngTokPerSec: Double = 0
|
| 201 |
+
var gotFirstToken = false
|
| 202 |
+
var inThink = false
|
| 203 |
+
|
| 204 |
+
let stream = engine.generate(prompt: formatted, maxTokens: maxTokensPerRun)
|
| 205 |
+
|
| 206 |
+
for await token in stream {
|
| 207 |
+
if !gotFirstToken {
|
| 208 |
+
firstTokenTime = CFAbsoluteTimeGetCurrent()
|
| 209 |
+
gotFirstToken = true
|
| 210 |
+
}
|
| 211 |
+
tokenCount += 1
|
| 212 |
+
lastEngTokPerSec = token.tokensPerSecond
|
| 213 |
+
|
| 214 |
+
// Clean special tokens + strip thinking
|
| 215 |
+
var clean = token.text
|
| 216 |
+
.replacingOccurrences(of: "<|im_end|>", with: "")
|
| 217 |
+
.replacingOccurrences(of: "<|im_start|>", with: "")
|
| 218 |
+
.replacingOccurrences(of: "<|endoftext|>", with: "")
|
| 219 |
+
if clean.contains("<think>") { inThink = true; clean = clean.replacingOccurrences(of: "<think>", with: "") }
|
| 220 |
+
if clean.contains("</think>") { inThink = false; clean = clean.replacingOccurrences(of: "</think>", with: "") }
|
| 221 |
+
if !inThink { output += clean }
|
| 222 |
+
}
|
| 223 |
+
|
| 224 |
+
let endTime = CFAbsoluteTimeGetCurrent()
|
| 225 |
+
let totalMs = (endTime - startTime) * 1000
|
| 226 |
+
let ttftMs = gotFirstToken ? (firstTokenTime - startTime) * 1000 : 0
|
| 227 |
+
|
| 228 |
+
// Calculate actual decode tok/s (exclude first token = prefill)
|
| 229 |
+
let decodeTokens = max(tokenCount - 1, 0)
|
| 230 |
+
let decodeTimeMs = gotFirstToken ? (endTime - firstTokenTime) * 1000 : 0
|
| 231 |
+
let decodeTokPerSec = decodeTimeMs > 0 ? Double(decodeTokens) / (decodeTimeMs / 1000.0) : 0
|
| 232 |
+
|
| 233 |
+
// Quality check
|
| 234 |
+
let hasGibberish = output.contains("!!!!") || output.contains("????")
|
| 235 |
+
let isEmpty = output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
|
| 236 |
+
let quality = hasGibberish ? "GIBBERISH" : isEmpty ? "EMPTY" : "OK"
|
| 237 |
+
|
| 238 |
+
// Estimate prefill tokens from TTFT log (rough)
|
| 239 |
+
let wordCount = prompt.text.split(separator: " ").count
|
| 240 |
+
let estimatedPrefillTokens = wordCount + 15 // +15 for chat template + system prompt tokens
|
| 241 |
+
|
| 242 |
+
let snippet = String(output.prefix(100))
|
| 243 |
+
|
| 244 |
+
return BenchmarkResult(
|
| 245 |
+
configName: config.name,
|
| 246 |
+
promptLabel: prompt.label,
|
| 247 |
+
promptText: prompt.text,
|
| 248 |
+
prefillTokens: estimatedPrefillTokens,
|
| 249 |
+
ttftMs: ttftMs,
|
| 250 |
+
tokensGenerated: tokenCount,
|
| 251 |
+
decodeTokPerSec: decodeTokPerSec,
|
| 252 |
+
engineTokPerSec: lastEngTokPerSec,
|
| 253 |
+
totalMs: totalMs,
|
| 254 |
+
outputSnippet: snippet,
|
| 255 |
+
quality: quality,
|
| 256 |
+
thermalState: thermal
|
| 257 |
+
)
|
| 258 |
+
}
|
| 259 |
+
|
| 260 |
+
var fullTSV: String {
|
| 261 |
+
var lines = "Config\tPrompt\tPrefillToks\tTTFT_ms\tGenToks\tDecodeTok/s\tEngTok/s\tTotalMs\tThermal\tQuality\tOutput\n"
|
| 262 |
+
for r in results {
|
| 263 |
+
lines += r.tsvLine + "\n"
|
| 264 |
+
}
|
| 265 |
+
return lines
|
| 266 |
+
}
|
| 267 |
+
}
|
| 268 |
+
|
| 269 |
+
// MARK: - Benchmark View
|
| 270 |
+
|
| 271 |
+
struct BenchmarkView: View {
|
| 272 |
+
@Environment(FlashMoEEngine.self) private var engine
|
| 273 |
+
let modelPath: String
|
| 274 |
+
|
| 275 |
+
@State private var runner = BenchmarkRunner()
|
| 276 |
+
@State private var selectedConfigs: Set<String> = Set(BenchmarkConfig.builtIn.map(\.id))
|
| 277 |
+
@State private var selectedPrompts: Set<String> = Set(BenchmarkPrompt.builtIn.map(\.id))
|
| 278 |
+
@State private var showSetup = true
|
| 279 |
+
|
| 280 |
+
var body: some View {
|
| 281 |
+
VStack(spacing: 0) {
|
| 282 |
+
// Status header
|
| 283 |
+
statusHeader
|
| 284 |
+
|
| 285 |
+
Divider()
|
| 286 |
+
|
| 287 |
+
// Results
|
| 288 |
+
if runner.results.isEmpty && !runner.isRunning {
|
| 289 |
+
emptyState
|
| 290 |
+
} else {
|
| 291 |
+
resultsList
|
| 292 |
+
}
|
| 293 |
+
}
|
| 294 |
+
.navigationTitle("Benchmark")
|
| 295 |
+
#if os(iOS)
|
| 296 |
+
.navigationBarTitleDisplayMode(.inline)
|
| 297 |
+
#endif
|
| 298 |
+
.toolbar {
|
| 299 |
+
ToolbarItem(placement: .confirmationAction) {
|
| 300 |
+
if runner.isRunning {
|
| 301 |
+
Button("Stop") {
|
| 302 |
+
engine.cancel()
|
| 303 |
+
runner.cancel()
|
| 304 |
+
}
|
| 305 |
+
.foregroundStyle(.red)
|
| 306 |
+
} else {
|
| 307 |
+
Button("Copy Results") {
|
| 308 |
+
copyResults()
|
| 309 |
+
}
|
| 310 |
+
.disabled(runner.results.isEmpty)
|
| 311 |
+
}
|
| 312 |
+
}
|
| 313 |
+
}
|
| 314 |
+
}
|
| 315 |
+
|
| 316 |
+
// MARK: - Status Header
|
| 317 |
+
|
| 318 |
+
private var statusHeader: some View {
|
| 319 |
+
VStack(spacing: 12) {
|
| 320 |
+
if runner.isRunning {
|
| 321 |
+
VStack(spacing: 6) {
|
| 322 |
+
ProgressView(
|
| 323 |
+
value: Double(runner.currentConfigIndex * runner.totalPrompts + runner.currentPromptIndex),
|
| 324 |
+
total: Double(runner.totalConfigs * runner.totalPrompts)
|
| 325 |
+
)
|
| 326 |
+
|
| 327 |
+
Text(runner.statusMessage)
|
| 328 |
+
.font(.headline)
|
| 329 |
+
|
| 330 |
+
if !runner.progressDetail.isEmpty {
|
| 331 |
+
Text(runner.progressDetail)
|
| 332 |
+
.font(.caption)
|
| 333 |
+
.foregroundStyle(.secondary)
|
| 334 |
+
}
|
| 335 |
+
}
|
| 336 |
+
.padding()
|
| 337 |
+
} else {
|
| 338 |
+
VStack(spacing: 8) {
|
| 339 |
+
Text(runner.statusMessage)
|
| 340 |
+
.font(.headline)
|
| 341 |
+
|
| 342 |
+
// Toggle setup panel
|
| 343 |
+
Button {
|
| 344 |
+
withAnimation { showSetup.toggle() }
|
| 345 |
+
} label: {
|
| 346 |
+
HStack {
|
| 347 |
+
Text("\(selectedConfigs.count) configs × \(selectedPrompts.count) prompts = \(selectedConfigs.count * selectedPrompts.count) runs")
|
| 348 |
+
.font(.caption)
|
| 349 |
+
.foregroundStyle(.secondary)
|
| 350 |
+
Image(systemName: showSetup ? "chevron.up" : "chevron.down")
|
| 351 |
+
.font(.caption2)
|
| 352 |
+
.foregroundStyle(.secondary)
|
| 353 |
+
}
|
| 354 |
+
}
|
| 355 |
+
.buttonStyle(.plain)
|
| 356 |
+
|
| 357 |
+
if showSetup {
|
| 358 |
+
VStack(alignment: .leading, spacing: 8) {
|
| 359 |
+
Text("Configs").font(.caption).fontWeight(.bold)
|
| 360 |
+
ForEach(BenchmarkConfig.builtIn) { cfg in
|
| 361 |
+
Toggle(cfg.name, isOn: Binding(
|
| 362 |
+
get: { selectedConfigs.contains(cfg.id) },
|
| 363 |
+
set: { if $0 { selectedConfigs.insert(cfg.id) } else { selectedConfigs.remove(cfg.id) } }
|
| 364 |
+
))
|
| 365 |
+
.font(.caption)
|
| 366 |
+
.toggleStyle(.switch)
|
| 367 |
+
}
|
| 368 |
+
|
| 369 |
+
Divider()
|
| 370 |
+
|
| 371 |
+
Text("Prompts").font(.caption).fontWeight(.bold)
|
| 372 |
+
ForEach(BenchmarkPrompt.builtIn) { p in
|
| 373 |
+
Toggle(p.label, isOn: Binding(
|
| 374 |
+
get: { selectedPrompts.contains(p.id) },
|
| 375 |
+
set: { if $0 { selectedPrompts.insert(p.id) } else { selectedPrompts.remove(p.id) } }
|
| 376 |
+
))
|
| 377 |
+
.font(.caption)
|
| 378 |
+
.toggleStyle(.switch)
|
| 379 |
+
}
|
| 380 |
+
}
|
| 381 |
+
.padding(.horizontal)
|
| 382 |
+
}
|
| 383 |
+
|
| 384 |
+
Button {
|
| 385 |
+
showSetup = false
|
| 386 |
+
Task {
|
| 387 |
+
await runner.run(engine: engine, modelPath: modelPath,
|
| 388 |
+
configIds: selectedConfigs, promptIds: selectedPrompts)
|
| 389 |
+
}
|
| 390 |
+
} label: {
|
| 391 |
+
Label("Run Benchmark", systemImage: "play.fill")
|
| 392 |
+
.frame(maxWidth: .infinity)
|
| 393 |
+
}
|
| 394 |
+
.buttonStyle(.borderedProminent)
|
| 395 |
+
.tint(.orange)
|
| 396 |
+
.disabled((engine.state != .ready && !runner.isRunning) || selectedConfigs.isEmpty || selectedPrompts.isEmpty)
|
| 397 |
+
}
|
| 398 |
+
.padding()
|
| 399 |
+
}
|
| 400 |
+
}
|
| 401 |
+
.background(.ultraThinMaterial)
|
| 402 |
+
}
|
| 403 |
+
|
| 404 |
+
// MARK: - Empty State
|
| 405 |
+
|
| 406 |
+
private var emptyState: some View {
|
| 407 |
+
VStack(spacing: 16) {
|
| 408 |
+
Spacer()
|
| 409 |
+
Image(systemName: "gauge.with.dots.needle.50percent")
|
| 410 |
+
.font(.system(size: 48))
|
| 411 |
+
.foregroundStyle(.secondary)
|
| 412 |
+
Text("Tap Run Benchmark to start")
|
| 413 |
+
.font(.headline)
|
| 414 |
+
.foregroundStyle(.secondary)
|
| 415 |
+
Text("Tests 5 prompts across 5 configurations.\nModel reloads between config changes.")
|
| 416 |
+
.font(.caption)
|
| 417 |
+
.foregroundStyle(.tertiary)
|
| 418 |
+
.multilineTextAlignment(.center)
|
| 419 |
+
Spacer()
|
| 420 |
+
}
|
| 421 |
+
.padding()
|
| 422 |
+
}
|
| 423 |
+
|
| 424 |
+
// MARK: - Results List
|
| 425 |
+
|
| 426 |
+
private var resultsList: some View {
|
| 427 |
+
ScrollView {
|
| 428 |
+
LazyVStack(spacing: 10) {
|
| 429 |
+
ForEach(runner.results) { result in
|
| 430 |
+
resultCard(result)
|
| 431 |
+
}
|
| 432 |
+
}
|
| 433 |
+
.padding(.horizontal, 12)
|
| 434 |
+
.padding(.vertical, 8)
|
| 435 |
+
}
|
| 436 |
+
}
|
| 437 |
+
|
| 438 |
+
private func resultCard(_ result: BenchmarkResult) -> some View {
|
| 439 |
+
VStack(alignment: .leading, spacing: 10) {
|
| 440 |
+
// Header: config + prompt + quality badge
|
| 441 |
+
HStack {
|
| 442 |
+
VStack(alignment: .leading, spacing: 2) {
|
| 443 |
+
Text(result.configName)
|
| 444 |
+
.font(.subheadline)
|
| 445 |
+
.fontWeight(.semibold)
|
| 446 |
+
Text(result.promptLabel)
|
| 447 |
+
.font(.caption)
|
| 448 |
+
.foregroundStyle(.secondary)
|
| 449 |
+
}
|
| 450 |
+
|
| 451 |
+
Spacer()
|
| 452 |
+
|
| 453 |
+
// Quality badge
|
| 454 |
+
Text(result.quality)
|
| 455 |
+
.font(.caption2)
|
| 456 |
+
.fontWeight(.bold)
|
| 457 |
+
.padding(.horizontal, 8)
|
| 458 |
+
.padding(.vertical, 4)
|
| 459 |
+
.background(qualityColor(result.quality))
|
| 460 |
+
.foregroundStyle(.white)
|
| 461 |
+
.clipShape(Capsule())
|
| 462 |
+
|
| 463 |
+
// Thermal indicator
|
| 464 |
+
if result.thermalState != "nominal" {
|
| 465 |
+
Image(systemName: thermalIcon(result.thermalState))
|
| 466 |
+
.font(.caption)
|
| 467 |
+
.foregroundStyle(thermalColor(result.thermalState))
|
| 468 |
+
}
|
| 469 |
+
}
|
| 470 |
+
|
| 471 |
+
// Metrics grid — 2 rows of 3
|
| 472 |
+
HStack(spacing: 0) {
|
| 473 |
+
metricTile("TTFT", String(format: "%.0f ms", result.ttftMs))
|
| 474 |
+
metricTile("Decode", String(format: "%.1f tok/s", result.decodeTokPerSec))
|
| 475 |
+
metricTile("Engine", String(format: "%.1f tok/s", result.engineTokPerSec))
|
| 476 |
+
}
|
| 477 |
+
|
| 478 |
+
HStack(spacing: 0) {
|
| 479 |
+
metricTile("Tokens", "\(result.tokensGenerated)")
|
| 480 |
+
metricTile("Prefill", "~\(result.prefillTokens) tok")
|
| 481 |
+
metricTile("Total", String(format: "%.1f s", result.totalMs / 1000))
|
| 482 |
+
}
|
| 483 |
+
|
| 484 |
+
// Output preview
|
| 485 |
+
if !result.outputSnippet.isEmpty {
|
| 486 |
+
Text(result.outputSnippet)
|
| 487 |
+
.font(.caption2)
|
| 488 |
+
.foregroundStyle(.secondary)
|
| 489 |
+
.lineLimit(2)
|
| 490 |
+
.padding(.top, 2)
|
| 491 |
+
}
|
| 492 |
+
}
|
| 493 |
+
.padding(12)
|
| 494 |
+
.background(Color(.secondarySystemGroupedBackground))
|
| 495 |
+
.clipShape(RoundedRectangle(cornerRadius: 12))
|
| 496 |
+
}
|
| 497 |
+
|
| 498 |
+
private func metricTile(_ label: String, _ value: String) -> some View {
|
| 499 |
+
VStack(spacing: 2) {
|
| 500 |
+
Text(value)
|
| 501 |
+
.font(.system(.callout, design: .rounded))
|
| 502 |
+
.fontWeight(.semibold)
|
| 503 |
+
Text(label)
|
| 504 |
+
.font(.caption2)
|
| 505 |
+
.foregroundStyle(.secondary)
|
| 506 |
+
}
|
| 507 |
+
.frame(maxWidth: .infinity)
|
| 508 |
+
}
|
| 509 |
+
|
| 510 |
+
private func qualityColor(_ quality: String) -> Color {
|
| 511 |
+
switch quality {
|
| 512 |
+
case "OK": return .green
|
| 513 |
+
case "GIBBERISH": return .red
|
| 514 |
+
case "EMPTY": return .orange
|
| 515 |
+
default: return .gray
|
| 516 |
+
}
|
| 517 |
+
}
|
| 518 |
+
|
| 519 |
+
private func thermalIcon(_ state: String) -> String {
|
| 520 |
+
switch state {
|
| 521 |
+
case "fair": return "thermometer.medium"
|
| 522 |
+
case "serious": return "thermometer.high"
|
| 523 |
+
case "critical": return "flame.fill"
|
| 524 |
+
default: return "thermometer.low"
|
| 525 |
+
}
|
| 526 |
+
}
|
| 527 |
+
|
| 528 |
+
private func thermalColor(_ state: String) -> Color {
|
| 529 |
+
switch state {
|
| 530 |
+
case "fair": return .yellow
|
| 531 |
+
case "serious": return .orange
|
| 532 |
+
case "critical": return .red
|
| 533 |
+
default: return .green
|
| 534 |
+
}
|
| 535 |
+
}
|
| 536 |
+
|
| 537 |
+
// MARK: - Actions
|
| 538 |
+
|
| 539 |
+
private func copyResults() {
|
| 540 |
+
#if os(iOS)
|
| 541 |
+
UIPasteboard.general.string = runner.fullTSV
|
| 542 |
+
#elseif os(macOS)
|
| 543 |
+
NSPasteboard.general.clearContents()
|
| 544 |
+
NSPasteboard.general.setString(runner.fullTSV, forType: .string)
|
| 545 |
+
#endif
|
| 546 |
+
}
|
| 547 |
+
}
|
FlashMoE-iOS/Views/ChatView.swift
ADDED
|
@@ -0,0 +1,538 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* ChatView.swift — Chat interface for Flash-MoE inference
|
| 3 |
+
*
|
| 4 |
+
* Streaming token display, stats overlay, conversation history.
|
| 5 |
+
*/
|
| 6 |
+
|
| 7 |
+
import SwiftUI
|
| 8 |
+
|
| 9 |
+
// MARK: - Chat Message Model
|
| 10 |
+
|
| 11 |
+
struct ChatMessage: Identifiable {
|
| 12 |
+
let id = UUID()
|
| 13 |
+
let role: Role
|
| 14 |
+
var text: String
|
| 15 |
+
let timestamp: Date
|
| 16 |
+
|
| 17 |
+
enum Role {
|
| 18 |
+
case user
|
| 19 |
+
case assistant
|
| 20 |
+
}
|
| 21 |
+
}
|
| 22 |
+
|
| 23 |
+
// MARK: - Chat View
|
| 24 |
+
|
| 25 |
+
struct ChatView: View {
|
| 26 |
+
@Environment(FlashMoEEngine.self) private var engine
|
| 27 |
+
@State private var messages: [ChatMessage] = []
|
| 28 |
+
@State private var inputText = ""
|
| 29 |
+
@State private var isGenerating = false
|
| 30 |
+
@State private var showStats = false
|
| 31 |
+
@State private var showModelInfo = false
|
| 32 |
+
@State private var showProfiler = false
|
| 33 |
+
@State private var showBenchmark = false
|
| 34 |
+
@State private var scrollAnchor = UUID() // updates on each token to trigger scroll
|
| 35 |
+
@FocusState private var inputFocused: Bool
|
| 36 |
+
@AppStorage("thinkingEnabled") private var thinkingEnabled: Bool = true
|
| 37 |
+
|
| 38 |
+
var body: some View {
|
| 39 |
+
VStack(spacing: 0) {
|
| 40 |
+
// Messages
|
| 41 |
+
ScrollViewReader { proxy in
|
| 42 |
+
ScrollView {
|
| 43 |
+
LazyVStack(alignment: .leading, spacing: 12) {
|
| 44 |
+
ForEach(messages) { message in
|
| 45 |
+
MessageBubble(message: message)
|
| 46 |
+
.id(message.id)
|
| 47 |
+
}
|
| 48 |
+
}
|
| 49 |
+
.padding()
|
| 50 |
+
}
|
| 51 |
+
.onTapGesture { inputFocused = false }
|
| 52 |
+
.onChange(of: messages.count) {
|
| 53 |
+
if let last = messages.last {
|
| 54 |
+
withAnimation(.easeOut(duration: 0.15)) {
|
| 55 |
+
proxy.scrollTo(last.id, anchor: .bottom)
|
| 56 |
+
}
|
| 57 |
+
}
|
| 58 |
+
}
|
| 59 |
+
.onChange(of: scrollAnchor) {
|
| 60 |
+
if let last = messages.last {
|
| 61 |
+
proxy.scrollTo(last.id, anchor: .bottom)
|
| 62 |
+
}
|
| 63 |
+
}
|
| 64 |
+
}
|
| 65 |
+
|
| 66 |
+
// Stats bar
|
| 67 |
+
if isGenerating || engine.tokensGenerated > 0 {
|
| 68 |
+
StatsBar(
|
| 69 |
+
tokensPerSecond: engine.tokensPerSecond,
|
| 70 |
+
tokensGenerated: engine.tokensGenerated,
|
| 71 |
+
isGenerating: isGenerating,
|
| 72 |
+
ttftMs: engine.timeToFirstToken
|
| 73 |
+
)
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
Divider()
|
| 77 |
+
|
| 78 |
+
// Input bar
|
| 79 |
+
HStack(spacing: 12) {
|
| 80 |
+
TextField("Message...", text: $inputText, axis: .vertical)
|
| 81 |
+
.textFieldStyle(.plain)
|
| 82 |
+
.lineLimit(1...5)
|
| 83 |
+
.focused($inputFocused)
|
| 84 |
+
.onSubmit { sendMessage() }
|
| 85 |
+
.disabled(isGenerating)
|
| 86 |
+
|
| 87 |
+
if isGenerating {
|
| 88 |
+
Button(action: { engine.cancel() }) {
|
| 89 |
+
Image(systemName: "stop.circle.fill")
|
| 90 |
+
.font(.title2)
|
| 91 |
+
.foregroundStyle(.red)
|
| 92 |
+
}
|
| 93 |
+
} else {
|
| 94 |
+
Button(action: sendMessage) {
|
| 95 |
+
Image(systemName: "arrow.up.circle.fill")
|
| 96 |
+
.font(.title2)
|
| 97 |
+
.foregroundStyle(inputText.isEmpty ? .gray : .blue)
|
| 98 |
+
}
|
| 99 |
+
.disabled(inputText.isEmpty)
|
| 100 |
+
}
|
| 101 |
+
}
|
| 102 |
+
.padding(.horizontal)
|
| 103 |
+
.padding(.vertical, 8)
|
| 104 |
+
}
|
| 105 |
+
.overlay(alignment: .bottom) {
|
| 106 |
+
if showProfiler {
|
| 107 |
+
ProfilerView(engine: engine)
|
| 108 |
+
.padding(.bottom, 80)
|
| 109 |
+
.transition(.move(edge: .bottom).combined(with: .opacity))
|
| 110 |
+
}
|
| 111 |
+
}
|
| 112 |
+
.animation(.easeInOut(duration: 0.25), value: showProfiler)
|
| 113 |
+
.navigationTitle("Flash-MoE")
|
| 114 |
+
#if os(iOS)
|
| 115 |
+
.navigationBarTitleDisplayMode(.inline)
|
| 116 |
+
.toolbar {
|
| 117 |
+
ToolbarItem(placement: .topBarLeading) {
|
| 118 |
+
Button(action: { showModelInfo = true }) {
|
| 119 |
+
Image(systemName: "cpu")
|
| 120 |
+
}
|
| 121 |
+
}
|
| 122 |
+
ToolbarItem(placement: .topBarTrailing) {
|
| 123 |
+
Menu {
|
| 124 |
+
Button("New Chat", systemImage: "plus.message") {
|
| 125 |
+
messages.removeAll()
|
| 126 |
+
engine.reset()
|
| 127 |
+
}
|
| 128 |
+
Button(showProfiler ? "Hide Profiler" : "Profiler", systemImage: "gauge.with.dots.needle.50percent") {
|
| 129 |
+
showProfiler.toggle()
|
| 130 |
+
}
|
| 131 |
+
Button("Show Stats", systemImage: "chart.bar") {
|
| 132 |
+
showStats.toggle()
|
| 133 |
+
}
|
| 134 |
+
Button("Benchmark", systemImage: "gauge.with.dots.needle.50percent") {
|
| 135 |
+
showBenchmark = true
|
| 136 |
+
}
|
| 137 |
+
Divider()
|
| 138 |
+
Button("Models & Settings", systemImage: "gearshape") {
|
| 139 |
+
messages.removeAll()
|
| 140 |
+
engine.reset()
|
| 141 |
+
engine.unloadModel()
|
| 142 |
+
}
|
| 143 |
+
} label: {
|
| 144 |
+
Image(systemName: "ellipsis.circle")
|
| 145 |
+
}
|
| 146 |
+
}
|
| 147 |
+
}
|
| 148 |
+
#else
|
| 149 |
+
.toolbar {
|
| 150 |
+
ToolbarItem(placement: .automatic) {
|
| 151 |
+
Button(action: { showModelInfo = true }) {
|
| 152 |
+
Image(systemName: "cpu")
|
| 153 |
+
}
|
| 154 |
+
}
|
| 155 |
+
ToolbarItem(placement: .automatic) {
|
| 156 |
+
Menu {
|
| 157 |
+
Button("New Chat", systemImage: "plus.message") {
|
| 158 |
+
messages.removeAll()
|
| 159 |
+
engine.reset()
|
| 160 |
+
}
|
| 161 |
+
Button(showProfiler ? "Hide Profiler" : "Profiler", systemImage: "gauge.with.dots.needle.50percent") {
|
| 162 |
+
showProfiler.toggle()
|
| 163 |
+
}
|
| 164 |
+
Button("Show Stats", systemImage: "chart.bar") {
|
| 165 |
+
showStats.toggle()
|
| 166 |
+
}
|
| 167 |
+
Button("Benchmark", systemImage: "gauge.with.dots.needle.50percent") {
|
| 168 |
+
showBenchmark = true
|
| 169 |
+
}
|
| 170 |
+
Divider()
|
| 171 |
+
Button("Models & Settings", systemImage: "gearshape") {
|
| 172 |
+
messages.removeAll()
|
| 173 |
+
engine.reset()
|
| 174 |
+
engine.unloadModel()
|
| 175 |
+
}
|
| 176 |
+
} label: {
|
| 177 |
+
Image(systemName: "ellipsis.circle")
|
| 178 |
+
}
|
| 179 |
+
}
|
| 180 |
+
}
|
| 181 |
+
#endif
|
| 182 |
+
.sheet(isPresented: $showModelInfo) {
|
| 183 |
+
ModelInfoSheet(info: engine.modelInfo)
|
| 184 |
+
}
|
| 185 |
+
.fullScreenCover(isPresented: $showBenchmark) {
|
| 186 |
+
NavigationStack {
|
| 187 |
+
BenchmarkView(modelPath: UserDefaults.standard.string(forKey: "lastLoadedModelPath") ?? "")
|
| 188 |
+
.environment(engine)
|
| 189 |
+
.toolbar {
|
| 190 |
+
ToolbarItem(placement: .cancellationAction) {
|
| 191 |
+
Button("Close") { showBenchmark = false }
|
| 192 |
+
}
|
| 193 |
+
}
|
| 194 |
+
}
|
| 195 |
+
}
|
| 196 |
+
}
|
| 197 |
+
|
| 198 |
+
private func sendMessage() {
|
| 199 |
+
let text = inputText.trimmingCharacters(in: .whitespacesAndNewlines)
|
| 200 |
+
guard !text.isEmpty else { return }
|
| 201 |
+
|
| 202 |
+
inputText = ""
|
| 203 |
+
let userMessage = ChatMessage(role: .user, text: text, timestamp: Date())
|
| 204 |
+
messages.append(userMessage)
|
| 205 |
+
|
| 206 |
+
// Start generation
|
| 207 |
+
isGenerating = true
|
| 208 |
+
let assistantMessage = ChatMessage(role: .assistant, text: "", timestamp: Date())
|
| 209 |
+
messages.append(assistantMessage)
|
| 210 |
+
let assistantIndex = messages.count - 1
|
| 211 |
+
|
| 212 |
+
Task {
|
| 213 |
+
let stream: AsyncStream<GenerationToken>
|
| 214 |
+
|
| 215 |
+
if engine.canContinue {
|
| 216 |
+
// Reuse KV cache — only process the new user turn
|
| 217 |
+
stream = engine.generateContinuation(userMessage: text, maxTokens: 2048)
|
| 218 |
+
} else {
|
| 219 |
+
// First message — full chat template with system prompt
|
| 220 |
+
let formattedPrompt = buildChatPrompt(userMessage: text)
|
| 221 |
+
stream = engine.generate(prompt: formattedPrompt, maxTokens: 2048)
|
| 222 |
+
}
|
| 223 |
+
|
| 224 |
+
var gotTokens = false
|
| 225 |
+
var tokenCount = 0
|
| 226 |
+
var inThinkBlock = false
|
| 227 |
+
for await token in stream {
|
| 228 |
+
gotTokens = true
|
| 229 |
+
tokenCount += 1
|
| 230 |
+
// Strip special tokens that leak through
|
| 231 |
+
var clean = token.text
|
| 232 |
+
.replacingOccurrences(of: "<|im_end|>", with: "")
|
| 233 |
+
.replacingOccurrences(of: "<|im_start|>", with: "")
|
| 234 |
+
.replacingOccurrences(of: "<|endoftext|>", with: "")
|
| 235 |
+
|
| 236 |
+
// When thinking is OFF, suppress <think>...</think> content
|
| 237 |
+
if !thinkingEnabled {
|
| 238 |
+
if clean.contains("<think>") {
|
| 239 |
+
inThinkBlock = true
|
| 240 |
+
clean = clean.replacingOccurrences(of: "<think>", with: "")
|
| 241 |
+
}
|
| 242 |
+
if clean.contains("</think>") {
|
| 243 |
+
inThinkBlock = false
|
| 244 |
+
clean = clean.replacingOccurrences(of: "</think>", with: "")
|
| 245 |
+
}
|
| 246 |
+
if inThinkBlock { clean = "" }
|
| 247 |
+
}
|
| 248 |
+
|
| 249 |
+
if !clean.isEmpty {
|
| 250 |
+
messages[assistantIndex].text += clean
|
| 251 |
+
// Auto-scroll every few tokens (not every token to avoid scroll jank)
|
| 252 |
+
if tokenCount % 3 == 0 {
|
| 253 |
+
scrollAnchor = UUID()
|
| 254 |
+
}
|
| 255 |
+
}
|
| 256 |
+
}
|
| 257 |
+
// Final scroll to catch last tokens
|
| 258 |
+
scrollAnchor = UUID()
|
| 259 |
+
|
| 260 |
+
// If continuation returned empty (context full), fall back to full generate
|
| 261 |
+
if !gotTokens && engine.canContinue {
|
| 262 |
+
engine.reset()
|
| 263 |
+
let formattedPrompt = buildChatPrompt(userMessage: text)
|
| 264 |
+
let fallbackStream = engine.generate(prompt: formattedPrompt, maxTokens: 2048)
|
| 265 |
+
tokenCount = 0
|
| 266 |
+
inThinkBlock = false
|
| 267 |
+
for await token in fallbackStream {
|
| 268 |
+
tokenCount += 1
|
| 269 |
+
var clean = token.text
|
| 270 |
+
.replacingOccurrences(of: "<|im_end|>", with: "")
|
| 271 |
+
.replacingOccurrences(of: "<|im_start|>", with: "")
|
| 272 |
+
.replacingOccurrences(of: "<|endoftext|>", with: "")
|
| 273 |
+
if !thinkingEnabled {
|
| 274 |
+
if clean.contains("<think>") { inThinkBlock = true; clean = clean.replacingOccurrences(of: "<think>", with: "") }
|
| 275 |
+
if clean.contains("</think>") { inThinkBlock = false; clean = clean.replacingOccurrences(of: "</think>", with: "") }
|
| 276 |
+
if inThinkBlock { clean = "" }
|
| 277 |
+
}
|
| 278 |
+
if !clean.isEmpty {
|
| 279 |
+
messages[assistantIndex].text += clean
|
| 280 |
+
if tokenCount % 3 == 0 {
|
| 281 |
+
scrollAnchor = UUID()
|
| 282 |
+
}
|
| 283 |
+
}
|
| 284 |
+
}
|
| 285 |
+
scrollAnchor = UUID()
|
| 286 |
+
}
|
| 287 |
+
|
| 288 |
+
isGenerating = false
|
| 289 |
+
}
|
| 290 |
+
}
|
| 291 |
+
|
| 292 |
+
/// Format conversation as Qwen chat template
|
| 293 |
+
private func buildChatPrompt(userMessage: String) -> String {
|
| 294 |
+
let sysContent = thinkingEnabled
|
| 295 |
+
? "You are a helpful assistant."
|
| 296 |
+
: "You are a helpful assistant. /no_think"
|
| 297 |
+
var prompt = "<|im_start|>system\n\(sysContent)<|im_end|>\n"
|
| 298 |
+
|
| 299 |
+
// Include conversation history (skip the empty assistant message we just appended)
|
| 300 |
+
for msg in messages.dropLast() {
|
| 301 |
+
switch msg.role {
|
| 302 |
+
case .user:
|
| 303 |
+
prompt += "<|im_start|>user\n\(msg.text)<|im_end|>\n"
|
| 304 |
+
case .assistant:
|
| 305 |
+
if !msg.text.isEmpty {
|
| 306 |
+
prompt += "<|im_start|>assistant\n\(msg.text)<|im_end|>\n"
|
| 307 |
+
}
|
| 308 |
+
}
|
| 309 |
+
}
|
| 310 |
+
|
| 311 |
+
// Qwen3.5 ALWAYS expects <think> after assistant start tag.
|
| 312 |
+
// Without it, the model emits EOS immediately. The /no_think
|
| 313 |
+
// instruction in the system prompt makes thinking brief, but
|
| 314 |
+
// the tag must be present for the model to generate anything.
|
| 315 |
+
prompt += "<|im_start|>assistant\n<think>\n"
|
| 316 |
+
return prompt
|
| 317 |
+
}
|
| 318 |
+
}
|
| 319 |
+
|
| 320 |
+
// MARK: - Message Bubble
|
| 321 |
+
|
| 322 |
+
struct MessageBubble: View {
|
| 323 |
+
let message: ChatMessage
|
| 324 |
+
@State private var showThinking = false
|
| 325 |
+
|
| 326 |
+
/// Split text into visible reply and thinking content
|
| 327 |
+
private var parsedContent: (think: String?, reply: String) {
|
| 328 |
+
let text = message.text
|
| 329 |
+
// Match <think>...</think> blocks
|
| 330 |
+
guard let thinkStart = text.range(of: "<think>"),
|
| 331 |
+
let thinkEnd = text.range(of: "</think>") else {
|
| 332 |
+
// No complete think block — check if still streaming thinking
|
| 333 |
+
if text.hasPrefix("<think>") {
|
| 334 |
+
let thinkBody = String(text.dropFirst("<think>".count))
|
| 335 |
+
return (think: thinkBody, reply: "")
|
| 336 |
+
}
|
| 337 |
+
return (think: nil, reply: text)
|
| 338 |
+
}
|
| 339 |
+
let thinkBody = String(text[thinkStart.upperBound..<thinkEnd.lowerBound]).trimmingCharacters(in: .whitespacesAndNewlines)
|
| 340 |
+
let reply = String(text[thinkEnd.upperBound...]).trimmingCharacters(in: .whitespacesAndNewlines)
|
| 341 |
+
return (think: thinkBody, reply: reply)
|
| 342 |
+
}
|
| 343 |
+
|
| 344 |
+
var body: some View {
|
| 345 |
+
HStack {
|
| 346 |
+
if message.role == .user { Spacer(minLength: 60) }
|
| 347 |
+
|
| 348 |
+
VStack(alignment: message.role == .user ? .trailing : .leading, spacing: 4) {
|
| 349 |
+
// Thinking disclosure (assistant only)
|
| 350 |
+
if message.role == .assistant, let thinkText = parsedContent.think, !thinkText.isEmpty {
|
| 351 |
+
DisclosureGroup(isExpanded: $showThinking) {
|
| 352 |
+
Text(thinkText)
|
| 353 |
+
.font(.caption)
|
| 354 |
+
.foregroundStyle(.secondary)
|
| 355 |
+
.textSelection(.enabled)
|
| 356 |
+
.padding(.horizontal, 10)
|
| 357 |
+
.padding(.vertical, 6)
|
| 358 |
+
} label: {
|
| 359 |
+
Label("Thinking...", systemImage: "brain")
|
| 360 |
+
.font(.caption)
|
| 361 |
+
.foregroundStyle(.secondary)
|
| 362 |
+
}
|
| 363 |
+
.padding(.horizontal, 14)
|
| 364 |
+
.padding(.vertical, 6)
|
| 365 |
+
#if os(iOS)
|
| 366 |
+
.background(Color(.systemGray6))
|
| 367 |
+
#else
|
| 368 |
+
.background(.thinMaterial)
|
| 369 |
+
#endif
|
| 370 |
+
.clipShape(RoundedRectangle(cornerRadius: 14))
|
| 371 |
+
}
|
| 372 |
+
|
| 373 |
+
// Main message text
|
| 374 |
+
let displayText = message.role == .assistant ? parsedContent.reply : message.text
|
| 375 |
+
if !displayText.isEmpty {
|
| 376 |
+
Text(displayText)
|
| 377 |
+
.textSelection(.enabled)
|
| 378 |
+
.padding(.horizontal, 14)
|
| 379 |
+
.padding(.vertical, 10)
|
| 380 |
+
#if os(iOS)
|
| 381 |
+
.background(message.role == .user ? Color.blue : Color(.systemGray5))
|
| 382 |
+
#else
|
| 383 |
+
.background(message.role == .user ? Color.blue : Color.secondary)
|
| 384 |
+
#endif
|
| 385 |
+
.foregroundStyle(message.role == .user ? .white : .primary)
|
| 386 |
+
.clipShape(RoundedRectangle(cornerRadius: 18))
|
| 387 |
+
}
|
| 388 |
+
|
| 389 |
+
if message.text.isEmpty && message.role == .assistant {
|
| 390 |
+
ThinkingIndicator()
|
| 391 |
+
.padding(.horizontal, 14)
|
| 392 |
+
.padding(.vertical, 10)
|
| 393 |
+
#if os(iOS)
|
| 394 |
+
.background(Color(.systemGray5))
|
| 395 |
+
#else
|
| 396 |
+
.background(.quaternary)
|
| 397 |
+
#endif
|
| 398 |
+
.clipShape(RoundedRectangle(cornerRadius: 18))
|
| 399 |
+
}
|
| 400 |
+
}
|
| 401 |
+
|
| 402 |
+
if message.role == .assistant { Spacer(minLength: 60) }
|
| 403 |
+
}
|
| 404 |
+
}
|
| 405 |
+
}
|
| 406 |
+
|
| 407 |
+
// MARK: - Thinking Indicator
|
| 408 |
+
|
| 409 |
+
struct ThinkingIndicator: View {
|
| 410 |
+
@State private var dotCount = 0
|
| 411 |
+
private let timer = Timer.publish(every: 0.4, on: .main, in: .common).autoconnect()
|
| 412 |
+
private let frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
|
| 413 |
+
|
| 414 |
+
var body: some View {
|
| 415 |
+
HStack(spacing: 6) {
|
| 416 |
+
Text(frames[dotCount % frames.count])
|
| 417 |
+
.font(.system(.body, design: .monospaced))
|
| 418 |
+
.foregroundStyle(.secondary)
|
| 419 |
+
Text("Thinking")
|
| 420 |
+
.font(.system(.body, design: .monospaced))
|
| 421 |
+
.foregroundStyle(.secondary)
|
| 422 |
+
}
|
| 423 |
+
.onReceive(timer) { _ in
|
| 424 |
+
dotCount += 1
|
| 425 |
+
}
|
| 426 |
+
}
|
| 427 |
+
}
|
| 428 |
+
|
| 429 |
+
// MARK: - Stats Bar
|
| 430 |
+
|
| 431 |
+
struct StatsBar: View {
|
| 432 |
+
let tokensPerSecond: Double
|
| 433 |
+
let tokensGenerated: Int
|
| 434 |
+
let isGenerating: Bool
|
| 435 |
+
var ttftMs: Double = 0
|
| 436 |
+
|
| 437 |
+
private var ttftText: String {
|
| 438 |
+
if ttftMs <= 0 { return "" }
|
| 439 |
+
if ttftMs > 500_000 {
|
| 440 |
+
return String(format: "%.1f min", ttftMs / 60_000)
|
| 441 |
+
} else if ttftMs > 1000 {
|
| 442 |
+
return String(format: "%.1fs", ttftMs / 1000)
|
| 443 |
+
} else {
|
| 444 |
+
return String(format: "%.0fms", ttftMs)
|
| 445 |
+
}
|
| 446 |
+
}
|
| 447 |
+
|
| 448 |
+
private var thermalLabel: String {
|
| 449 |
+
switch ProcessInfo.processInfo.thermalState {
|
| 450 |
+
case .nominal: return "\u{1F7E2} Cool"
|
| 451 |
+
case .fair: return "\u{1F7E1} Warm"
|
| 452 |
+
case .serious: return "\u{1F7E0} Hot"
|
| 453 |
+
case .critical: return "\u{1F534} Critical"
|
| 454 |
+
@unknown default: return "\u{2753} Unknown"
|
| 455 |
+
}
|
| 456 |
+
}
|
| 457 |
+
|
| 458 |
+
var body: some View {
|
| 459 |
+
HStack(spacing: 12) {
|
| 460 |
+
Label(String(format: "%.1f tok/s", tokensPerSecond), systemImage: "speedometer")
|
| 461 |
+
.font(.caption)
|
| 462 |
+
.foregroundStyle(.secondary)
|
| 463 |
+
|
| 464 |
+
Label("\(tokensGenerated) tokens", systemImage: "number")
|
| 465 |
+
.font(.caption)
|
| 466 |
+
.foregroundStyle(.secondary)
|
| 467 |
+
|
| 468 |
+
if !ttftText.isEmpty {
|
| 469 |
+
Label(ttftText, systemImage: "clock")
|
| 470 |
+
.font(.caption)
|
| 471 |
+
.foregroundStyle(.secondary)
|
| 472 |
+
}
|
| 473 |
+
|
| 474 |
+
Text(thermalLabel)
|
| 475 |
+
.font(.caption)
|
| 476 |
+
.foregroundStyle(.secondary)
|
| 477 |
+
|
| 478 |
+
Spacer()
|
| 479 |
+
|
| 480 |
+
if isGenerating {
|
| 481 |
+
ProgressView()
|
| 482 |
+
.scaleEffect(0.7)
|
| 483 |
+
}
|
| 484 |
+
}
|
| 485 |
+
.padding(.horizontal)
|
| 486 |
+
.padding(.vertical, 6)
|
| 487 |
+
.background(.ultraThinMaterial)
|
| 488 |
+
}
|
| 489 |
+
}
|
| 490 |
+
|
| 491 |
+
// MARK: - Model Info Sheet
|
| 492 |
+
|
| 493 |
+
struct ModelInfoSheet: View {
|
| 494 |
+
let info: ModelInfo?
|
| 495 |
+
@Environment(\.dismiss) private var dismiss
|
| 496 |
+
|
| 497 |
+
var body: some View {
|
| 498 |
+
NavigationStack {
|
| 499 |
+
if let info {
|
| 500 |
+
List {
|
| 501 |
+
Section("Architecture") {
|
| 502 |
+
InfoRow(label: "Layers", value: "\(info.numLayers)")
|
| 503 |
+
InfoRow(label: "Experts", value: "\(info.numExperts) (K=\(info.activeExpertsK))")
|
| 504 |
+
InfoRow(label: "Hidden Dim", value: "\(info.hiddenDim)")
|
| 505 |
+
InfoRow(label: "Vocab Size", value: "\(info.vocabSize)")
|
| 506 |
+
}
|
| 507 |
+
Section("Storage") {
|
| 508 |
+
InfoRow(label: "Weights", value: String(format: "%.1f MB", info.weightFileMB))
|
| 509 |
+
InfoRow(label: "Experts", value: String(format: "%.1f MB", info.expertFileMB))
|
| 510 |
+
InfoRow(label: "Total", value: String(format: "%.1f GB", info.totalSizeMB / 1024))
|
| 511 |
+
}
|
| 512 |
+
}
|
| 513 |
+
.navigationTitle("Model Info")
|
| 514 |
+
} else {
|
| 515 |
+
Text("No model loaded")
|
| 516 |
+
}
|
| 517 |
+
}
|
| 518 |
+
#if os(iOS)
|
| 519 |
+
.presentationDetents([.medium])
|
| 520 |
+
#endif
|
| 521 |
+
}
|
| 522 |
+
}
|
| 523 |
+
|
| 524 |
+
struct InfoRow: View {
|
| 525 |
+
let label: String
|
| 526 |
+
let value: String
|
| 527 |
+
|
| 528 |
+
var body: some View {
|
| 529 |
+
HStack {
|
| 530 |
+
Text(label)
|
| 531 |
+
.foregroundStyle(.secondary)
|
| 532 |
+
Spacer()
|
| 533 |
+
Text(value)
|
| 534 |
+
.fontDesign(.monospaced)
|
| 535 |
+
}
|
| 536 |
+
}
|
| 537 |
+
}
|
| 538 |
+
|
FlashMoE-iOS/Views/ContentView.swift
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* ContentView.swift — Root navigation view
|
| 3 |
+
*
|
| 4 |
+
* Shows model list if no model loaded, chat if model is ready.
|
| 5 |
+
*/
|
| 6 |
+
|
| 7 |
+
import SwiftUI
|
| 8 |
+
|
| 9 |
+
struct ContentView: View {
|
| 10 |
+
@Environment(FlashMoEEngine.self) private var engine
|
| 11 |
+
|
| 12 |
+
var body: some View {
|
| 13 |
+
NavigationStack {
|
| 14 |
+
switch engine.state {
|
| 15 |
+
case .idle, .loading, .error:
|
| 16 |
+
ModelListView()
|
| 17 |
+
case .ready, .generating:
|
| 18 |
+
ChatView()
|
| 19 |
+
}
|
| 20 |
+
}
|
| 21 |
+
}
|
| 22 |
+
}
|
FlashMoE-iOS/Views/ModelDownloadRow.swift
ADDED
|
@@ -0,0 +1,232 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* ModelDownloadRow.swift — Download progress UI for catalog entries
|
| 3 |
+
*
|
| 4 |
+
* Shows different states: available for download, downloading with progress,
|
| 5 |
+
* paused, failed with retry, or completed.
|
| 6 |
+
*/
|
| 7 |
+
|
| 8 |
+
import SwiftUI
|
| 9 |
+
|
| 10 |
+
struct ModelDownloadRow: View {
|
| 11 |
+
let entry: CatalogEntry
|
| 12 |
+
let downloadManager: DownloadManager
|
| 13 |
+
let isDownloaded: Bool
|
| 14 |
+
|
| 15 |
+
private var quantColor: Color {
|
| 16 |
+
switch entry.quantization.lowercased() {
|
| 17 |
+
case "4-bit": return .blue
|
| 18 |
+
case "2-bit": return .orange
|
| 19 |
+
case "tiered": return .purple
|
| 20 |
+
default: return .gray
|
| 21 |
+
}
|
| 22 |
+
}
|
| 23 |
+
|
| 24 |
+
private var isActiveDownload: Bool {
|
| 25 |
+
downloadManager.activeDownload?.catalogId == entry.id
|
| 26 |
+
}
|
| 27 |
+
|
| 28 |
+
private var downloadStatus: DownloadStatus? {
|
| 29 |
+
guard isActiveDownload else { return nil }
|
| 30 |
+
return downloadManager.activeDownload?.status
|
| 31 |
+
}
|
| 32 |
+
|
| 33 |
+
var body: some View {
|
| 34 |
+
VStack(alignment: .leading, spacing: 8) {
|
| 35 |
+
// Header row
|
| 36 |
+
HStack {
|
| 37 |
+
VStack(alignment: .leading, spacing: 2) {
|
| 38 |
+
Text(entry.displayName)
|
| 39 |
+
.font(.headline)
|
| 40 |
+
Text(entry.description)
|
| 41 |
+
.font(.caption)
|
| 42 |
+
.foregroundStyle(.secondary)
|
| 43 |
+
.lineLimit(2)
|
| 44 |
+
}
|
| 45 |
+
|
| 46 |
+
Spacer()
|
| 47 |
+
|
| 48 |
+
QuantBadge(text: entry.quantization, color: quantColor)
|
| 49 |
+
}
|
| 50 |
+
|
| 51 |
+
// Status-specific content
|
| 52 |
+
if isDownloaded {
|
| 53 |
+
downloadedView
|
| 54 |
+
} else if let status = downloadStatus {
|
| 55 |
+
switch status {
|
| 56 |
+
case .downloading:
|
| 57 |
+
downloadingView
|
| 58 |
+
case .paused:
|
| 59 |
+
pausedView
|
| 60 |
+
case .failed:
|
| 61 |
+
failedView
|
| 62 |
+
case .complete:
|
| 63 |
+
downloadedView
|
| 64 |
+
}
|
| 65 |
+
} else {
|
| 66 |
+
availableView
|
| 67 |
+
}
|
| 68 |
+
}
|
| 69 |
+
.padding(.vertical, 4)
|
| 70 |
+
}
|
| 71 |
+
|
| 72 |
+
// MARK: - State Views
|
| 73 |
+
|
| 74 |
+
private var availableView: some View {
|
| 75 |
+
HStack {
|
| 76 |
+
Label(formatSize(entry.totalSizeBytes), systemImage: "internaldrive")
|
| 77 |
+
.font(.caption)
|
| 78 |
+
.foregroundStyle(.secondary)
|
| 79 |
+
|
| 80 |
+
Spacer()
|
| 81 |
+
|
| 82 |
+
Button {
|
| 83 |
+
downloadManager.startDownload(entry: entry)
|
| 84 |
+
} label: {
|
| 85 |
+
Label("Download", systemImage: "arrow.down.circle.fill")
|
| 86 |
+
.font(.subheadline.weight(.medium))
|
| 87 |
+
}
|
| 88 |
+
.buttonStyle(.borderedProminent)
|
| 89 |
+
.tint(.blue)
|
| 90 |
+
}
|
| 91 |
+
}
|
| 92 |
+
|
| 93 |
+
private var downloadingView: some View {
|
| 94 |
+
VStack(alignment: .leading, spacing: 6) {
|
| 95 |
+
ProgressView(value: downloadManager.overallProgress)
|
| 96 |
+
.tint(.blue)
|
| 97 |
+
|
| 98 |
+
HStack {
|
| 99 |
+
// Progress text
|
| 100 |
+
Text("\(formatSize(downloadManager.bytesDownloaded)) / \(formatSize(downloadManager.totalBytes))")
|
| 101 |
+
.font(.caption)
|
| 102 |
+
.foregroundStyle(.secondary)
|
| 103 |
+
|
| 104 |
+
if downloadManager.downloadSpeed > 0 {
|
| 105 |
+
Text("(\(formatSpeed(downloadManager.downloadSpeed)))")
|
| 106 |
+
.font(.caption)
|
| 107 |
+
.foregroundStyle(.secondary)
|
| 108 |
+
}
|
| 109 |
+
|
| 110 |
+
Spacer()
|
| 111 |
+
|
| 112 |
+
// Current file
|
| 113 |
+
if let currentFile = downloadManager.activeDownload?.currentFile {
|
| 114 |
+
Text(shortFilename(currentFile))
|
| 115 |
+
.font(.caption2)
|
| 116 |
+
.foregroundStyle(.tertiary)
|
| 117 |
+
}
|
| 118 |
+
|
| 119 |
+
// Pause button
|
| 120 |
+
Button {
|
| 121 |
+
downloadManager.pauseDownload()
|
| 122 |
+
} label: {
|
| 123 |
+
Image(systemName: "pause.circle.fill")
|
| 124 |
+
.foregroundStyle(.orange)
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
// Cancel button
|
| 128 |
+
Button {
|
| 129 |
+
downloadManager.cancelDownload()
|
| 130 |
+
} label: {
|
| 131 |
+
Image(systemName: "xmark.circle.fill")
|
| 132 |
+
.foregroundStyle(.red)
|
| 133 |
+
}
|
| 134 |
+
}
|
| 135 |
+
}
|
| 136 |
+
}
|
| 137 |
+
|
| 138 |
+
private var pausedView: some View {
|
| 139 |
+
VStack(alignment: .leading, spacing: 6) {
|
| 140 |
+
ProgressView(value: downloadManager.overallProgress)
|
| 141 |
+
.tint(.orange)
|
| 142 |
+
|
| 143 |
+
HStack {
|
| 144 |
+
Text("Paused — \(formatSize(downloadManager.bytesDownloaded)) / \(formatSize(downloadManager.totalBytes))")
|
| 145 |
+
.font(.caption)
|
| 146 |
+
.foregroundStyle(.orange)
|
| 147 |
+
|
| 148 |
+
Spacer()
|
| 149 |
+
|
| 150 |
+
Button {
|
| 151 |
+
downloadManager.resumeDownload()
|
| 152 |
+
} label: {
|
| 153 |
+
Label("Resume", systemImage: "play.circle.fill")
|
| 154 |
+
.font(.subheadline.weight(.medium))
|
| 155 |
+
}
|
| 156 |
+
.buttonStyle(.borderedProminent)
|
| 157 |
+
.tint(.orange)
|
| 158 |
+
|
| 159 |
+
Button {
|
| 160 |
+
downloadManager.cancelDownload()
|
| 161 |
+
} label: {
|
| 162 |
+
Image(systemName: "xmark.circle.fill")
|
| 163 |
+
.foregroundStyle(.red)
|
| 164 |
+
}
|
| 165 |
+
}
|
| 166 |
+
}
|
| 167 |
+
}
|
| 168 |
+
|
| 169 |
+
private var failedView: some View {
|
| 170 |
+
VStack(alignment: .leading, spacing: 6) {
|
| 171 |
+
if let error = downloadManager.activeDownload?.errorMessage ?? downloadManager.error {
|
| 172 |
+
Label(error, systemImage: "exclamationmark.triangle.fill")
|
| 173 |
+
.font(.caption)
|
| 174 |
+
.foregroundStyle(.red)
|
| 175 |
+
}
|
| 176 |
+
|
| 177 |
+
HStack {
|
| 178 |
+
Text("\(formatSize(downloadManager.bytesDownloaded)) downloaded")
|
| 179 |
+
.font(.caption)
|
| 180 |
+
.foregroundStyle(.secondary)
|
| 181 |
+
|
| 182 |
+
Spacer()
|
| 183 |
+
|
| 184 |
+
Button {
|
| 185 |
+
downloadManager.resumeDownload()
|
| 186 |
+
} label: {
|
| 187 |
+
Label("Retry", systemImage: "arrow.clockwise.circle.fill")
|
| 188 |
+
.font(.subheadline.weight(.medium))
|
| 189 |
+
}
|
| 190 |
+
.buttonStyle(.borderedProminent)
|
| 191 |
+
.tint(.red)
|
| 192 |
+
|
| 193 |
+
Button {
|
| 194 |
+
downloadManager.cancelDownload()
|
| 195 |
+
} label: {
|
| 196 |
+
Image(systemName: "xmark.circle.fill")
|
| 197 |
+
.foregroundStyle(.secondary)
|
| 198 |
+
}
|
| 199 |
+
}
|
| 200 |
+
}
|
| 201 |
+
}
|
| 202 |
+
|
| 203 |
+
private var downloadedView: some View {
|
| 204 |
+
HStack {
|
| 205 |
+
Label("Downloaded", systemImage: "checkmark.circle.fill")
|
| 206 |
+
.font(.caption)
|
| 207 |
+
.foregroundStyle(.green)
|
| 208 |
+
|
| 209 |
+
Spacer()
|
| 210 |
+
}
|
| 211 |
+
}
|
| 212 |
+
|
| 213 |
+
// MARK: - Formatting
|
| 214 |
+
|
| 215 |
+
private func formatSize(_ bytes: UInt64) -> String {
|
| 216 |
+
let gb = Double(bytes) / (1024.0 * 1024.0 * 1024.0)
|
| 217 |
+
if gb >= 1 { return String(format: "%.1f GB", gb) }
|
| 218 |
+
let mb = Double(bytes) / (1024.0 * 1024.0)
|
| 219 |
+
return String(format: "%.0f MB", mb)
|
| 220 |
+
}
|
| 221 |
+
|
| 222 |
+
private func formatSpeed(_ bytesPerSec: Double) -> String {
|
| 223 |
+
let mbps = bytesPerSec / (1024 * 1024)
|
| 224 |
+
return String(format: "%.1f MB/s", mbps)
|
| 225 |
+
}
|
| 226 |
+
|
| 227 |
+
private func shortFilename(_ filename: String) -> String {
|
| 228 |
+
(filename as NSString).lastPathComponent
|
| 229 |
+
}
|
| 230 |
+
}
|
| 231 |
+
|
| 232 |
+
// QuantBadge is defined in ModelListView.swift
|
FlashMoE-iOS/Views/ModelListView.swift
ADDED
|
@@ -0,0 +1,1212 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* ModelListView.swift — Model discovery and loading
|
| 3 |
+
*
|
| 4 |
+
* Lists locally available models and allows downloading from HuggingFace.
|
| 5 |
+
* For v1, supports loading models already present on device.
|
| 6 |
+
*/
|
| 7 |
+
|
| 8 |
+
import SwiftUI
|
| 9 |
+
import UniformTypeIdentifiers
|
| 10 |
+
#if os(macOS)
|
| 11 |
+
import AppKit
|
| 12 |
+
#endif
|
| 13 |
+
|
| 14 |
+
// MARK: - Local Model Entry
|
| 15 |
+
|
| 16 |
+
struct LocalModel: Identifiable {
|
| 17 |
+
let id = UUID()
|
| 18 |
+
let name: String
|
| 19 |
+
let path: String
|
| 20 |
+
let sizeBytes: UInt64
|
| 21 |
+
let hasTiered: Bool
|
| 22 |
+
let has4bit: Bool
|
| 23 |
+
let has2bit: Bool
|
| 24 |
+
|
| 25 |
+
var sizeMB: Double { Double(sizeBytes) / 1_048_576 }
|
| 26 |
+
var sizeGB: Double { sizeMB / 1024 }
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
// MARK: - Setting Info
|
| 30 |
+
|
| 31 |
+
struct SettingInfo: Identifiable {
|
| 32 |
+
let id: String
|
| 33 |
+
let title: String
|
| 34 |
+
let analogy: String
|
| 35 |
+
let technical: String
|
| 36 |
+
|
| 37 |
+
static let all: [String: SettingInfo] = [
|
| 38 |
+
"activeExperts": SettingInfo(
|
| 39 |
+
id: "activeExperts",
|
| 40 |
+
title: "Active Experts (K)",
|
| 41 |
+
analogy: "Imagine asking a question to a room of 256 specialists. K controls how many you consult. K=8 means you ask 8 experts and combine their answers. K=4 means you only ask 4 — faster (less reading from disk) but you might miss a specialist who had a great insight.",
|
| 42 |
+
technical: "Each transformer layer routes the token to K out of 256 experts via a learned gating network. Each expert is a ~1.7MB weight matrix loaded from SSD via pread(). Lower K = fewer SSD reads per layer = proportionally less I/O time (the dominant bottleneck at 56% of per-token latency). Quality degrades gracefully because the router still picks the best K from the full vocabulary."
|
| 43 |
+
),
|
| 44 |
+
"ioFanout": SettingInfo(
|
| 45 |
+
id: "ioFanout",
|
| 46 |
+
title: "Expert I/O Fanout",
|
| 47 |
+
analogy: "Think of reading a book page. Instead of reading the whole page in one go, you split it into strips and read them all simultaneously with multiple eyes. Fanout splits each expert weight file read into parallel chunks so the SSD controller can serve them concurrently.",
|
| 48 |
+
technical: "Each expert (~1.7MB for 35B) is read via pread(). Fanout splits this into N page-aligned chunks dispatched via GCD dispatch_group_async. NVMe controllers have multiple queues and can serve parallel reads faster than a single large read. Best value depends on expert size vs NVMe page size (4KB). Diminishing returns above 4 chunks."
|
| 49 |
+
),
|
| 50 |
+
"cmdMerge": SettingInfo(
|
| 51 |
+
id: "cmdMerge",
|
| 52 |
+
title: "CMD1+CMD2 Merge",
|
| 53 |
+
analogy: "Like combining two errands into one trip instead of driving home between them. CMD1 (attention projections) and CMD2 (output projection + normalization) are separate GPU tasks. Merging them avoids the roundtrip of 'submit, wait, create new, submit' for each of the 30 linear attention layers.",
|
| 54 |
+
technical: "For linear attention layers (GatedDeltaNet), the CPU phase between CMD1 and CMD2 is empty — the GPU already computed everything. CMD2's dispatches (o_proj matmul, residual add, RMS norm, routing, shared expert) are appended to CMD1 with pipeline barriers. Saves ~0.05-0.1ms per layer x 30 layers = 1.5-3ms per token."
|
| 55 |
+
),
|
| 56 |
+
"fusedAttention": SettingInfo(
|
| 57 |
+
id: "fusedAttention",
|
| 58 |
+
title: "Fused Attention",
|
| 59 |
+
analogy: "Standard attention is like a three-step cooking recipe: measure all ingredients (Q@K scores), mix them (softmax), then combine (scores@V). Fused attention does all three in one pass — like a skilled chef who seasons, mixes, and plates in a single flowing motion. Less cleanup between steps.",
|
| 60 |
+
technical: "Replaces 3 separate GPU kernel dispatches (attn_scores, attn_softmax, attn_values) with a single fused kernel using FlashAttention-2 online softmax. Processes KV positions in blocks of 64, maintaining running max/sum/output. Eliminates 2 command encoder transitions per full-attention layer (10 layers). Uses unnormalized accumulation with single final division."
|
| 61 |
+
),
|
| 62 |
+
"fusedExpert": SettingInfo(
|
| 63 |
+
id: "fusedExpert",
|
| 64 |
+
title: "Fused Expert Kernel",
|
| 65 |
+
analogy: "Each expert normally does three separate calculations: gate, up, and activation. It's like washing, drying, and folding laundry in three separate trips. The fused kernel does all three in one pass through the data — one trip, everything done.",
|
| 66 |
+
technical: "Combines gate_proj matmul + up_proj matmul + SiLU activation into a single Metal compute kernel (fused_gate_up_swiglu). Both gate and up dot products are computed in one loop over the input vector, then SiLU is applied immediately. Reduces from 3 GPU dispatches to 1 per expert, saving command encoder overhead for K experts x 40 layers."
|
| 67 |
+
),
|
| 68 |
+
"expertPrefetch": SettingInfo(
|
| 69 |
+
id: "expertPrefetch",
|
| 70 |
+
title: "Expert Prefetch",
|
| 71 |
+
analogy: "While the kitchen (GPU) is cooking layer 5's dish, the waiter (CPU) runs ahead to the pantry (SSD) to grab ingredients for layer 6. When the kitchen finishes layer 5, the ingredients for layer 6 are already on the counter — no waiting.",
|
| 72 |
+
technical: "After CMD3(N) is submitted (deferred GPU execution), the system predicts which experts layer N+1 will need based on routing history. Those experts are pread() into Set B buffers asynchronously. When layer N+1 reaches its I/O phase, prefetch hits skip the pread entirely. Misses fall through to normal loading. Overlaps ~2.4ms of I/O with GPU compute time."
|
| 73 |
+
),
|
| 74 |
+
"fp16Accum": SettingInfo(
|
| 75 |
+
id: "fp16Accum",
|
| 76 |
+
title: "FP16 Accumulation",
|
| 77 |
+
analogy: "Imagine counting coins on a kitchen scale that rounds to one decimal. Each coin adds a tiny rounding error. After 500 coins, you might be off by one. But the scale reads twice as fast. FP16 does math at 2x the speed of FP32, but accumulates small rounding errors over hundreds of additions.",
|
| 78 |
+
technical: "The dequant matvec inner loop changes from float32 to float16 accumulation. Apple's A-series GPU has dedicated fp16 ALUs at 2x throughput. The FMA becomes half-precision: fma(half(nibble), half(scale*x), half(bias*x)). Final output is promoted to float32 via simd_sum. Risk: fp16 has ~3 decimal digits; sums of 512+ elements may lose precision."
|
| 79 |
+
),
|
| 80 |
+
"fp8KV": SettingInfo(
|
| 81 |
+
id: "fp8KV",
|
| 82 |
+
title: "FP8 KV Cache",
|
| 83 |
+
analogy: "The KV cache is like a notebook where the model writes down what it's seen. FP32 uses a full page per note. FP8 uses a quarter page — same content, just more compressed handwriting. You fit 4x more notes in the same notebook, so the model can remember 4x more conversation.",
|
| 84 |
+
technical: "Stores attention Key and Value vectors in FP8 E4M3 format (1 byte vs 4 bytes per element) with per-position dynamic scaling. Encoding: absmax/240 scale factor, each float clipped and quantized to 8-bit (1 sign, 4 exponent, 3 mantissa). Decoding is inline during attention compute. 4x memory reduction enables 4x longer context at the same memory budget."
|
| 85 |
+
),
|
| 86 |
+
"maxContext": SettingInfo(
|
| 87 |
+
id: "maxContext",
|
| 88 |
+
title: "Max Context Length",
|
| 89 |
+
analogy: "Context length is how far back the model can 'see' in the conversation. Like a person's short-term memory — 4K tokens is the last few minutes, 32K is the last hour. More context = better understanding of the conversation, but uses more memory.",
|
| 90 |
+
technical: "Sets the maximum sequence length for KV cache allocation. Memory cost: num_full_attn_layers x 2 (K+V) x kv_heads x head_dim x bytes_per_elem x positions. For the 35B with 10 full-attn layers: 40KB/pos (FP32) or 10KB/pos (FP8). Auto mode uses os_proc_available_memory() to pick the largest safe value."
|
| 91 |
+
),
|
| 92 |
+
"slidingWindow": SettingInfo(
|
| 93 |
+
id: "slidingWindow",
|
| 94 |
+
title: "Sliding Window",
|
| 95 |
+
analogy: "Instead of remembering everything forever (which fills up memory), the full-attention layers only look at the last N tokens — like a window sliding along the conversation. But the 30 linear attention layers still remember everything through their state matrices. It's like having both short-term and long-term memory working together.",
|
| 96 |
+
technical: "Implements a circular KV buffer for full attention layers. Write: cache_pos = kv->len % window_size. Read: attend only to the most recent window_size positions. The 30 GatedDeltaNet layers maintain full context via their 128x128 state matrices (O(1) memory). Only the 10 full attention layers are windowed. With window 4096 + FP8: fixed 40MB KV regardless of conversation length."
|
| 97 |
+
),
|
| 98 |
+
"thinking": SettingInfo(
|
| 99 |
+
id: "thinking",
|
| 100 |
+
title: "Thinking Mode",
|
| 101 |
+
analogy: "Like a student who shows their work before giving the final answer. The model reasons step-by-step inside <think> tags before responding. This usually produces better answers, but takes more tokens (and time). At low K values, the model may get stuck thinking forever — disable it for speed.",
|
| 102 |
+
technical: "The chat template includes a <think> tag after the assistant turn header. The model generates reasoning tokens inside the think block, then emits </think> before the actual response. Think budget caps the maximum thinking tokens and force-emits </think>. Set to -1 to disable thinking entirely (removes <think> from the template)."
|
| 103 |
+
),
|
| 104 |
+
"h2oBudget": SettingInfo(
|
| 105 |
+
id: "h2oBudget",
|
| 106 |
+
title: "KV Cache Budget (H\u{2082}O)",
|
| 107 |
+
analogy: "Instead of remembering everything or only recent things, H\u{2082}O is like a smart student who keeps notes on the most important points from the whole lecture plus the last few minutes. It tracks which tokens the model pays most attention to and keeps those, dropping the rest.",
|
| 108 |
+
technical: "H\u{2082}O (Heavy Hitter Oracle) eviction keeps 3 categories: attention sink tokens (first 4, always high-attention), recent tokens (last N), and heavy hitters (highest cumulative attention score). When the cache exceeds the budget, the lowest-scored non-protected positions are evicted. Reduces KV memory by 50-70% with minimal quality loss. Based on 'H\u{2082}O: Heavy-Hitter Oracle' (Zhang et al., 2023)."
|
| 109 |
+
),
|
| 110 |
+
"ropeScaling": SettingInfo(
|
| 111 |
+
id: "ropeScaling",
|
| 112 |
+
title: "RoPE Scaling",
|
| 113 |
+
analogy: "The model learns position by spinning numbers at specific speeds. Beyond the training window, these numbers 'go off the map' and output degrades. RoPE scaling adjusts the spin speeds so longer conversations stay on-map. Linear is the simplest (just slow everything down). NTK-aware is smarter (adjusts different dimensions differently). YaRN is best (combines NTK with an attention temperature fix).",
|
| 114 |
+
technical: "Rotary Position Embeddings encode position as frequency-domain rotations: freq_i = 1/base^(2i/d). Linear scaling divides position by the scale factor. NTK-aware scales the base: base' = base * s^(d/(d-2)), spreading frequencies more evenly. YaRN adds per-dimension interpolation based on wavelength vs. original context, plus an attention logit temperature t = sqrt(0.1*ln(s)+1). References: NTK-aware (arXiv:2306.15595), YaRN (arXiv:2309.00071)."
|
| 115 |
+
),
|
| 116 |
+
"prefillBatch": SettingInfo(
|
| 117 |
+
id: "prefillBatch",
|
| 118 |
+
title: "Prefill Batch Size",
|
| 119 |
+
analogy: "Normally the model reads each prompt word one at a time. Batched prefill reads multiple words simultaneously, sharing the weight data across all of them. It is like a teacher reading one textbook page to the whole class at once instead of whispering it to each student individually.",
|
| 120 |
+
technical: "Converts per-token GEMV (matrix-vector multiply) into batched GEMM (matrix-matrix multiply) during the prefill phase. Each weight row is read once from SSD and multiplied against N input vectors simultaneously, amortizing I/O cost. Uses dedicated Metal kernels with per-token accumulators. Memory overhead: ~33 MB for pfb=32. Causal attention uses the FlashAttention-2 unnormalized accumulator pattern for numerical stability."
|
| 121 |
+
),
|
| 122 |
+
]
|
| 123 |
+
}
|
| 124 |
+
|
| 125 |
+
// MARK: - Model List View
|
| 126 |
+
|
| 127 |
+
struct ModelListView: View {
|
| 128 |
+
@Environment(FlashMoEEngine.self) private var engine
|
| 129 |
+
@State private var localModels: [LocalModel] = []
|
| 130 |
+
@State private var isScanning = true
|
| 131 |
+
@State private var loadError: String?
|
| 132 |
+
@State private var selectedModel: LocalModel?
|
| 133 |
+
|
| 134 |
+
// Build info — set at compile time via COMMIT_HASH build setting or fallback
|
| 135 |
+
private var buildCommitShort: String {
|
| 136 |
+
Bundle.main.infoDictionary?["CFBundleVersion"] as? String ?? "dev"
|
| 137 |
+
}
|
| 138 |
+
@AppStorage("cacheIOSplit") private var cacheIOSplit: Int = 1
|
| 139 |
+
@AppStorage("activeExpertsK") private var activeExpertsK: Int = 0
|
| 140 |
+
@AppStorage("cmdMergeEnabled") private var cmdMergeEnabled: Bool = true
|
| 141 |
+
@AppStorage("fusedAttention") private var fusedAttention: Bool = false
|
| 142 |
+
@AppStorage("thinkingEnabled") private var thinkingEnabled: Bool = true
|
| 143 |
+
@AppStorage("thinkBudget") private var thinkBudget: Int = 2048
|
| 144 |
+
@AppStorage("expertPrefetch") private var expertPrefetch: Bool = false
|
| 145 |
+
@AppStorage("fusedExpert") private var fusedExpert: Bool = true
|
| 146 |
+
@AppStorage("fp16Accumulation") private var fp16Accumulation: Bool = false
|
| 147 |
+
@AppStorage("fp8KVCache") private var fp8KVCache: Bool = false
|
| 148 |
+
@AppStorage("maxContext") private var maxContext: Int = 0
|
| 149 |
+
@AppStorage("slidingWindow") private var slidingWindow: Int = 0
|
| 150 |
+
@AppStorage("h2oBudget") private var h2oBudget: Int = 0
|
| 151 |
+
@AppStorage("ropeScaling") private var ropeScaling: Int = 0 // encoded: mode * 10 + factor_index
|
| 152 |
+
@AppStorage("prefillBatch") private var prefillBatch: Int = 1
|
| 153 |
+
@State private var settingInfo: SettingInfo? = nil
|
| 154 |
+
@State private var showFilePicker = false
|
| 155 |
+
@State private var modelToExport: LocalModel? = nil
|
| 156 |
+
@State private var importedBookmark: Data? = nil
|
| 157 |
+
@State private var showImportActionAlert = false
|
| 158 |
+
@State private var pendingImportURL: URL? = nil
|
| 159 |
+
@State private var importProgress: String? = nil
|
| 160 |
+
@State private var modelToDelete: LocalModel? = nil
|
| 161 |
+
@State private var customRepoURL: String = ""
|
| 162 |
+
@State private var showCustomURLError: String? = nil
|
| 163 |
+
@State private var showSpeed = false
|
| 164 |
+
@State private var showPipeline = false
|
| 165 |
+
@State private var showContext = false
|
| 166 |
+
@State private var showGeneration = false
|
| 167 |
+
@State private var customEntries: [CatalogEntry] = []
|
| 168 |
+
@State private var isResolvingURL = false
|
| 169 |
+
private let downloadManager = DownloadManager.shared
|
| 170 |
+
|
| 171 |
+
var body: some View {
|
| 172 |
+
List {
|
| 173 |
+
Section {
|
| 174 |
+
headerView
|
| 175 |
+
}
|
| 176 |
+
.listRowBackground(Color.clear)
|
| 177 |
+
|
| 178 |
+
if isScanning {
|
| 179 |
+
Section {
|
| 180 |
+
HStack {
|
| 181 |
+
ProgressView()
|
| 182 |
+
Text("Scanning for models...")
|
| 183 |
+
.foregroundStyle(.secondary)
|
| 184 |
+
}
|
| 185 |
+
}
|
| 186 |
+
} else if localModels.isEmpty {
|
| 187 |
+
Section {
|
| 188 |
+
VStack(alignment: .leading, spacing: 8) {
|
| 189 |
+
Text("No models found")
|
| 190 |
+
.font(.headline)
|
| 191 |
+
#if os(iOS)
|
| 192 |
+
Text("Download a model below, or transfer one via Files.app.")
|
| 193 |
+
.font(.caption)
|
| 194 |
+
.foregroundStyle(.secondary)
|
| 195 |
+
#else
|
| 196 |
+
Text("Download a model below, or use \"Open Model Folder\" to load from disk.")
|
| 197 |
+
.font(.caption)
|
| 198 |
+
.foregroundStyle(.secondary)
|
| 199 |
+
#endif
|
| 200 |
+
}
|
| 201 |
+
.padding(.vertical, 4)
|
| 202 |
+
}
|
| 203 |
+
} else {
|
| 204 |
+
Section("On Device") {
|
| 205 |
+
ForEach(localModels) { model in
|
| 206 |
+
ModelRow(model: model, isLoading: engine.state == .loading && selectedModel?.id == model.id)
|
| 207 |
+
.onTapGesture { loadModel(model) }
|
| 208 |
+
.swipeActions(edge: .leading) {
|
| 209 |
+
Button {
|
| 210 |
+
modelToExport = model
|
| 211 |
+
} label: {
|
| 212 |
+
Label("Move to Files", systemImage: "square.and.arrow.up")
|
| 213 |
+
}
|
| 214 |
+
.tint(.blue)
|
| 215 |
+
}
|
| 216 |
+
.swipeActions(edge: .trailing, allowsFullSwipe: false) {
|
| 217 |
+
Button(role: .destructive) {
|
| 218 |
+
modelToDelete = model
|
| 219 |
+
} label: {
|
| 220 |
+
Label("Delete", systemImage: "trash")
|
| 221 |
+
}
|
| 222 |
+
}
|
| 223 |
+
}
|
| 224 |
+
}
|
| 225 |
+
}
|
| 226 |
+
|
| 227 |
+
// Custom URL download
|
| 228 |
+
Section("Add Model from URL") {
|
| 229 |
+
HStack {
|
| 230 |
+
TextField("HuggingFace repo (user/model)", text: $customRepoURL)
|
| 231 |
+
.textFieldStyle(.roundedBorder)
|
| 232 |
+
.autocorrectionDisabled()
|
| 233 |
+
#if os(iOS)
|
| 234 |
+
.textInputAutocapitalization(.never)
|
| 235 |
+
.keyboardType(.URL)
|
| 236 |
+
#endif
|
| 237 |
+
Button {
|
| 238 |
+
startCustomDownload()
|
| 239 |
+
} label: {
|
| 240 |
+
if isResolvingURL {
|
| 241 |
+
ProgressView()
|
| 242 |
+
.controlSize(.small)
|
| 243 |
+
} else {
|
| 244 |
+
Image(systemName: "plus.circle.fill")
|
| 245 |
+
.font(.title3)
|
| 246 |
+
}
|
| 247 |
+
}
|
| 248 |
+
.disabled(customRepoURL.trimmingCharacters(in: .whitespaces).isEmpty || isResolvingURL)
|
| 249 |
+
}
|
| 250 |
+
if let error = showCustomURLError {
|
| 251 |
+
Text(error)
|
| 252 |
+
.font(.caption)
|
| 253 |
+
.foregroundStyle(.red)
|
| 254 |
+
}
|
| 255 |
+
Text("Enter a HuggingFace repo ID (e.g. alexintosh/Qwen3.5-35B-A3B-Q4-FlashMoE) or full URL.")
|
| 256 |
+
.font(.caption)
|
| 257 |
+
.foregroundStyle(.secondary)
|
| 258 |
+
}
|
| 259 |
+
|
| 260 |
+
// Download section — hide models already on device
|
| 261 |
+
let localNames = Set(localModels.map { $0.name.lowercased() })
|
| 262 |
+
let allEntries = ModelCatalog.models + customEntries
|
| 263 |
+
let availableEntries = allEntries.filter { entry in
|
| 264 |
+
// Hide if already downloaded on device
|
| 265 |
+
let isOnDevice = downloadManager.isModelDownloaded(entry.id) || localNames.contains(entry.id.lowercased())
|
| 266 |
+
let hasActiveDownload = downloadManager.activeDownload?.catalogId == entry.id
|
| 267 |
+
&& downloadManager.activeDownload?.status != .complete
|
| 268 |
+
return !isOnDevice || hasActiveDownload
|
| 269 |
+
}
|
| 270 |
+
if !availableEntries.isEmpty {
|
| 271 |
+
Section("Download from HuggingFace") {
|
| 272 |
+
ForEach(availableEntries) { entry in
|
| 273 |
+
let hasActiveDownload = downloadManager.activeDownload?.catalogId == entry.id
|
| 274 |
+
&& downloadManager.activeDownload?.status != .complete
|
| 275 |
+
ModelDownloadRow(
|
| 276 |
+
entry: entry,
|
| 277 |
+
downloadManager: downloadManager,
|
| 278 |
+
isDownloaded: !hasActiveDownload && downloadManager.isModelDownloaded(entry.id)
|
| 279 |
+
)
|
| 280 |
+
}
|
| 281 |
+
}
|
| 282 |
+
}
|
| 283 |
+
|
| 284 |
+
// ---- Speed Settings ----
|
| 285 |
+
Section {
|
| 286 |
+
DisclosureGroup("Speed", isExpanded: $showSpeed) {
|
| 287 |
+
Picker(selection: $activeExpertsK) {
|
| 288 |
+
Text("Model default").tag(0)
|
| 289 |
+
Text("K=2 (fastest)").tag(2)
|
| 290 |
+
Text("K=3").tag(3)
|
| 291 |
+
Text("K=4").tag(4)
|
| 292 |
+
Text("K=5").tag(5)
|
| 293 |
+
Text("K=6").tag(6)
|
| 294 |
+
Text("K=7").tag(7)
|
| 295 |
+
Text("K=8").tag(8)
|
| 296 |
+
Text("K=9").tag(9)
|
| 297 |
+
Text("K=10 (full)").tag(10)
|
| 298 |
+
} label: { settingLabel("Active Experts (K)", key: "activeExperts") }
|
| 299 |
+
.pickerStyle(.menu)
|
| 300 |
+
|
| 301 |
+
Picker(selection: $cacheIOSplit) {
|
| 302 |
+
Text("Off").tag(1)
|
| 303 |
+
Text("2 chunks").tag(2)
|
| 304 |
+
Text("4 chunks").tag(4)
|
| 305 |
+
Text("8 chunks").tag(8)
|
| 306 |
+
} label: { settingLabel("I/O Fanout", key: "ioFanout") }
|
| 307 |
+
.pickerStyle(.menu)
|
| 308 |
+
|
| 309 |
+
Picker(selection: $prefillBatch) {
|
| 310 |
+
Text("Off (1)").tag(1)
|
| 311 |
+
Text("8 tokens").tag(8)
|
| 312 |
+
Text("16 tokens").tag(16)
|
| 313 |
+
Text("32 tokens").tag(32)
|
| 314 |
+
} label: { settingLabel("Prefill Batch Size", key: "prefillBatch") }
|
| 315 |
+
.pickerStyle(.menu)
|
| 316 |
+
|
| 317 |
+
Toggle(isOn: $fp16Accumulation) { settingLabel("FP16 Accumulation", key: "fp16Accum") }
|
| 318 |
+
}
|
| 319 |
+
}
|
| 320 |
+
|
| 321 |
+
// ---- GPU Pipeline ----
|
| 322 |
+
Section {
|
| 323 |
+
DisclosureGroup("GPU Pipeline", isExpanded: $showPipeline) {
|
| 324 |
+
Toggle(isOn: $cmdMergeEnabled) { settingLabel("CMD1+CMD2 Merge", key: "cmdMerge") }
|
| 325 |
+
Toggle(isOn: $fusedAttention) { settingLabel("Fused Attention", key: "fusedAttention") }
|
| 326 |
+
Toggle(isOn: $fusedExpert) { settingLabel("Fused Expert Kernel", key: "fusedExpert") }
|
| 327 |
+
Toggle(isOn: $expertPrefetch) { settingLabel("Expert Prefetch", key: "expertPrefetch") }
|
| 328 |
+
}
|
| 329 |
+
}
|
| 330 |
+
|
| 331 |
+
// ---- Context & Memory ----
|
| 332 |
+
Section {
|
| 333 |
+
DisclosureGroup("Context & Memory", isExpanded: $showContext) {
|
| 334 |
+
Toggle(isOn: $fp8KVCache) { settingLabel("FP8 KV Cache", key: "fp8KV") }
|
| 335 |
+
|
| 336 |
+
Picker(selection: $maxContext) {
|
| 337 |
+
Text("Auto").tag(0)
|
| 338 |
+
Text("4K").tag(4096)
|
| 339 |
+
Text("8K").tag(8192)
|
| 340 |
+
Text("16K").tag(16384)
|
| 341 |
+
Text("32K").tag(32768)
|
| 342 |
+
} label: { settingLabel("Max Context", key: "maxContext") }
|
| 343 |
+
.pickerStyle(.menu)
|
| 344 |
+
|
| 345 |
+
Picker(selection: $slidingWindow) {
|
| 346 |
+
Text("Off").tag(0)
|
| 347 |
+
Text("2K").tag(2048)
|
| 348 |
+
Text("4K").tag(4096)
|
| 349 |
+
Text("8K").tag(8192)
|
| 350 |
+
} label: { settingLabel("Sliding Window", key: "slidingWindow") }
|
| 351 |
+
.pickerStyle(.menu)
|
| 352 |
+
|
| 353 |
+
Picker(selection: $h2oBudget) {
|
| 354 |
+
Text("Off").tag(0)
|
| 355 |
+
Text("256").tag(256)
|
| 356 |
+
Text("512").tag(512)
|
| 357 |
+
Text("1024").tag(1024)
|
| 358 |
+
Text("2048").tag(2048)
|
| 359 |
+
} label: { settingLabel("KV Cache Budget (H\u{2082}O)", key: "h2oBudget") }
|
| 360 |
+
.pickerStyle(.menu)
|
| 361 |
+
|
| 362 |
+
Picker(selection: $ropeScaling) {
|
| 363 |
+
Text("Off").tag(0)
|
| 364 |
+
Text("Linear 2x").tag(12)
|
| 365 |
+
Text("Linear 4x").tag(14)
|
| 366 |
+
Text("NTK-aware 2x").tag(22)
|
| 367 |
+
Text("NTK-aware 4x").tag(24)
|
| 368 |
+
Text("YaRN 2x").tag(32)
|
| 369 |
+
Text("YaRN 4x").tag(34)
|
| 370 |
+
} label: { settingLabel("RoPE Scaling", key: "ropeScaling") }
|
| 371 |
+
.pickerStyle(.menu)
|
| 372 |
+
}
|
| 373 |
+
}
|
| 374 |
+
|
| 375 |
+
// ---- Generation ----
|
| 376 |
+
Section {
|
| 377 |
+
DisclosureGroup("Generation", isExpanded: $showGeneration) {
|
| 378 |
+
Toggle(isOn: $thinkingEnabled) { settingLabel("Thinking", key: "thinking") }
|
| 379 |
+
|
| 380 |
+
if thinkingEnabled {
|
| 381 |
+
Picker("Think Budget", selection: $thinkBudget) {
|
| 382 |
+
Text("128").tag(128)
|
| 383 |
+
Text("256").tag(256)
|
| 384 |
+
Text("512").tag(512)
|
| 385 |
+
Text("1024").tag(1024)
|
| 386 |
+
Text("2048").tag(2048)
|
| 387 |
+
Text("Unlimited").tag(0)
|
| 388 |
+
}
|
| 389 |
+
.pickerStyle(.menu)
|
| 390 |
+
}
|
| 391 |
+
}
|
| 392 |
+
}
|
| 393 |
+
|
| 394 |
+
// ---- Apply ----
|
| 395 |
+
Section {
|
| 396 |
+
Button {
|
| 397 |
+
if let model = selectedModel ?? localModels.first {
|
| 398 |
+
Task {
|
| 399 |
+
if engine.state == .ready {
|
| 400 |
+
engine.unloadModel()
|
| 401 |
+
try? await Task.sleep(for: .milliseconds(300))
|
| 402 |
+
}
|
| 403 |
+
loadModel(model)
|
| 404 |
+
}
|
| 405 |
+
}
|
| 406 |
+
} label: {
|
| 407 |
+
Label("Reload Model", systemImage: "arrow.clockwise")
|
| 408 |
+
.frame(maxWidth: .infinity)
|
| 409 |
+
}
|
| 410 |
+
.buttonStyle(.borderedProminent)
|
| 411 |
+
.tint(.orange)
|
| 412 |
+
.disabled(localModels.isEmpty)
|
| 413 |
+
}
|
| 414 |
+
|
| 415 |
+
// Build info
|
| 416 |
+
Section {
|
| 417 |
+
Text("Branch: develop • \(buildCommitShort)")
|
| 418 |
+
.font(.caption2)
|
| 419 |
+
.foregroundStyle(.tertiary)
|
| 420 |
+
.frame(maxWidth: .infinity)
|
| 421 |
+
.listRowBackground(Color.clear)
|
| 422 |
+
}
|
| 423 |
+
|
| 424 |
+
if let error = downloadManager.error,
|
| 425 |
+
downloadManager.activeDownload == nil {
|
| 426 |
+
Section {
|
| 427 |
+
Label(error, systemImage: "exclamationmark.triangle")
|
| 428 |
+
.foregroundStyle(.red)
|
| 429 |
+
.font(.caption)
|
| 430 |
+
}
|
| 431 |
+
}
|
| 432 |
+
|
| 433 |
+
if case .error(let msg) = engine.state {
|
| 434 |
+
Section {
|
| 435 |
+
Label(msg, systemImage: "exclamationmark.triangle")
|
| 436 |
+
.foregroundStyle(.red)
|
| 437 |
+
.font(.caption)
|
| 438 |
+
}
|
| 439 |
+
}
|
| 440 |
+
}
|
| 441 |
+
.navigationTitle("Flash-MoE")
|
| 442 |
+
.toolbar {
|
| 443 |
+
#if os(iOS)
|
| 444 |
+
ToolbarItem(placement: .topBarTrailing) {
|
| 445 |
+
Button {
|
| 446 |
+
showFilePicker = true
|
| 447 |
+
} label: {
|
| 448 |
+
Label("Import", systemImage: "folder.badge.plus")
|
| 449 |
+
}
|
| 450 |
+
}
|
| 451 |
+
#else
|
| 452 |
+
ToolbarItem(placement: .automatic) {
|
| 453 |
+
Button {
|
| 454 |
+
openMacModelFolder()
|
| 455 |
+
} label: {
|
| 456 |
+
Label("Open Model Folder...", systemImage: "folder.badge.plus")
|
| 457 |
+
}
|
| 458 |
+
}
|
| 459 |
+
#endif
|
| 460 |
+
}
|
| 461 |
+
.sheet(item: $settingInfo) { info in
|
| 462 |
+
NavigationStack {
|
| 463 |
+
ScrollView {
|
| 464 |
+
VStack(alignment: .leading, spacing: 16) {
|
| 465 |
+
Text(info.analogy)
|
| 466 |
+
.font(.body)
|
| 467 |
+
.foregroundStyle(.primary)
|
| 468 |
+
|
| 469 |
+
Divider()
|
| 470 |
+
|
| 471 |
+
Text("How it works")
|
| 472 |
+
.font(.headline)
|
| 473 |
+
.foregroundStyle(.secondary)
|
| 474 |
+
|
| 475 |
+
Text(info.technical)
|
| 476 |
+
.font(.callout)
|
| 477 |
+
.foregroundStyle(.secondary)
|
| 478 |
+
}
|
| 479 |
+
.padding()
|
| 480 |
+
}
|
| 481 |
+
.navigationTitle(info.title)
|
| 482 |
+
#if os(iOS)
|
| 483 |
+
.navigationBarTitleDisplayMode(.inline)
|
| 484 |
+
#endif
|
| 485 |
+
.toolbar {
|
| 486 |
+
ToolbarItem(placement: .confirmationAction) {
|
| 487 |
+
Button("Done") { settingInfo = nil }
|
| 488 |
+
}
|
| 489 |
+
}
|
| 490 |
+
}
|
| 491 |
+
.presentationDetents([.medium])
|
| 492 |
+
}
|
| 493 |
+
#if os(iOS)
|
| 494 |
+
.sheet(isPresented: $showFilePicker) {
|
| 495 |
+
FolderImportPicker { url in
|
| 496 |
+
pendingImportURL = url
|
| 497 |
+
showImportActionAlert = true
|
| 498 |
+
}
|
| 499 |
+
}
|
| 500 |
+
#endif
|
| 501 |
+
#if os(iOS)
|
| 502 |
+
.alert("Import Model", isPresented: $showImportActionAlert) {
|
| 503 |
+
Button("Link (Bookmark)") {
|
| 504 |
+
if let url = pendingImportURL {
|
| 505 |
+
handleImportedFolder(url)
|
| 506 |
+
}
|
| 507 |
+
pendingImportURL = nil
|
| 508 |
+
}
|
| 509 |
+
Button("Move to App (Documents)") {
|
| 510 |
+
if let url = pendingImportURL {
|
| 511 |
+
moveImportedFolderToDocuments(url)
|
| 512 |
+
}
|
| 513 |
+
pendingImportURL = nil
|
| 514 |
+
}
|
| 515 |
+
Button("Cancel", role: .cancel) {
|
| 516 |
+
pendingImportURL = nil
|
| 517 |
+
}
|
| 518 |
+
} message: {
|
| 519 |
+
Text("Link keeps the model in its current location. Move to App copies it into the app's Documents folder for better reliability.")
|
| 520 |
+
}
|
| 521 |
+
#endif
|
| 522 |
+
.alert("Delete Model", isPresented: Binding(
|
| 523 |
+
get: { modelToDelete != nil },
|
| 524 |
+
set: { if !$0 { modelToDelete = nil } }
|
| 525 |
+
)) {
|
| 526 |
+
Button("Delete", role: .destructive) {
|
| 527 |
+
if let model = modelToDelete {
|
| 528 |
+
deleteModel(model)
|
| 529 |
+
}
|
| 530 |
+
modelToDelete = nil
|
| 531 |
+
}
|
| 532 |
+
Button("Cancel", role: .cancel) { modelToDelete = nil }
|
| 533 |
+
} message: {
|
| 534 |
+
Text("Delete \"\(modelToDelete?.name ?? "")\" (\(String(format: "%.1f GB", modelToDelete?.sizeGB ?? 0.0)))? This cannot be undone.")
|
| 535 |
+
}
|
| 536 |
+
#if os(iOS)
|
| 537 |
+
.sheet(item: $modelToExport) { model in
|
| 538 |
+
FolderExportPicker(sourceURL: URL(fileURLWithPath: model.path)) { destURL in
|
| 539 |
+
// moveToService already moved the files — just refresh the model list
|
| 540 |
+
print("[export] Model moved to: \(destURL.path)")
|
| 541 |
+
scanForModels()
|
| 542 |
+
}
|
| 543 |
+
}
|
| 544 |
+
#endif
|
| 545 |
+
.overlay {
|
| 546 |
+
if let progress = importProgress {
|
| 547 |
+
VStack(spacing: 12) {
|
| 548 |
+
ProgressView()
|
| 549 |
+
Text(progress)
|
| 550 |
+
.font(.caption)
|
| 551 |
+
.foregroundStyle(.secondary)
|
| 552 |
+
}
|
| 553 |
+
.padding(24)
|
| 554 |
+
.background(.ultraThinMaterial)
|
| 555 |
+
.clipShape(RoundedRectangle(cornerRadius: 12))
|
| 556 |
+
}
|
| 557 |
+
}
|
| 558 |
+
.onAppear {
|
| 559 |
+
scanForModels()
|
| 560 |
+
#if os(iOS)
|
| 561 |
+
restoreBookmarks()
|
| 562 |
+
#endif
|
| 563 |
+
}
|
| 564 |
+
.refreshable { scanForModels() }
|
| 565 |
+
.onChange(of: downloadManager.activeDownload?.status) { _, newStatus in
|
| 566 |
+
if newStatus == .complete {
|
| 567 |
+
scanForModels()
|
| 568 |
+
}
|
| 569 |
+
}
|
| 570 |
+
}
|
| 571 |
+
|
| 572 |
+
private var headerView: some View {
|
| 573 |
+
VStack(spacing: 8) {
|
| 574 |
+
Image(systemName: "bolt.fill")
|
| 575 |
+
.font(.system(size: 48))
|
| 576 |
+
.foregroundStyle(.orange)
|
| 577 |
+
Text("Flash-MoE")
|
| 578 |
+
.font(.largeTitle.bold())
|
| 579 |
+
#if os(iOS)
|
| 580 |
+
Text("Run massive MoE models on iPhone")
|
| 581 |
+
.font(.subheadline)
|
| 582 |
+
.foregroundStyle(.secondary)
|
| 583 |
+
#else
|
| 584 |
+
Text("Run massive MoE models on your Mac")
|
| 585 |
+
.font(.subheadline)
|
| 586 |
+
.foregroundStyle(.secondary)
|
| 587 |
+
#endif
|
| 588 |
+
}
|
| 589 |
+
.frame(maxWidth: .infinity)
|
| 590 |
+
.padding(.vertical)
|
| 591 |
+
}
|
| 592 |
+
|
| 593 |
+
private func scanForModels() {
|
| 594 |
+
isScanning = true
|
| 595 |
+
localModels = []
|
| 596 |
+
|
| 597 |
+
Task {
|
| 598 |
+
let models = await ModelScanner.scanLocalModels()
|
| 599 |
+
await MainActor.run {
|
| 600 |
+
localModels = models
|
| 601 |
+
isScanning = false
|
| 602 |
+
}
|
| 603 |
+
}
|
| 604 |
+
}
|
| 605 |
+
|
| 606 |
+
#if os(macOS)
|
| 607 |
+
private func openMacModelFolder() {
|
| 608 |
+
let panel = NSOpenPanel()
|
| 609 |
+
panel.canChooseDirectories = true
|
| 610 |
+
panel.canChooseFiles = false
|
| 611 |
+
panel.allowsMultipleSelection = false
|
| 612 |
+
panel.message = "Select a Flash-MoE model directory"
|
| 613 |
+
panel.prompt = "Open"
|
| 614 |
+
if panel.runModal() == .OK, let url = panel.url {
|
| 615 |
+
// On macOS, no sandbox bookmark needed for user-selected directories.
|
| 616 |
+
// Just store the path and scan.
|
| 617 |
+
let fm = FileManager.default
|
| 618 |
+
if FlashMoEEngine.validateModel(at: url.path) {
|
| 619 |
+
// Valid model directory — add to scanned paths
|
| 620 |
+
var paths = UserDefaults.standard.stringArray(forKey: "macModelPaths") ?? []
|
| 621 |
+
if !paths.contains(url.path) {
|
| 622 |
+
paths.append(url.path)
|
| 623 |
+
UserDefaults.standard.set(paths, forKey: "macModelPaths")
|
| 624 |
+
}
|
| 625 |
+
scanForModels()
|
| 626 |
+
} else {
|
| 627 |
+
// Check if it contains model subdirectories
|
| 628 |
+
var foundModel = false
|
| 629 |
+
if let entries = try? fm.contentsOfDirectory(atPath: url.path) {
|
| 630 |
+
for entry in entries {
|
| 631 |
+
let fullPath = (url.path as NSString).appendingPathComponent(entry)
|
| 632 |
+
if FlashMoEEngine.validateModel(at: fullPath) {
|
| 633 |
+
foundModel = true
|
| 634 |
+
break
|
| 635 |
+
}
|
| 636 |
+
}
|
| 637 |
+
}
|
| 638 |
+
var paths = UserDefaults.standard.stringArray(forKey: "macModelPaths") ?? []
|
| 639 |
+
if !paths.contains(url.path) {
|
| 640 |
+
paths.append(url.path)
|
| 641 |
+
UserDefaults.standard.set(paths, forKey: "macModelPaths")
|
| 642 |
+
}
|
| 643 |
+
scanForModels()
|
| 644 |
+
}
|
| 645 |
+
}
|
| 646 |
+
}
|
| 647 |
+
#endif
|
| 648 |
+
|
| 649 |
+
private func loadModel(_ model: LocalModel) {
|
| 650 |
+
guard engine.state != .loading && engine.state != .generating else { return }
|
| 651 |
+
selectedModel = model
|
| 652 |
+
// Store path for benchmark access
|
| 653 |
+
UserDefaults.standard.set(model.path, forKey: "lastLoadedModelPath")
|
| 654 |
+
|
| 655 |
+
// Use picker value, or auto-detect for 397B if user hasn't set a preference
|
| 656 |
+
let activeK: Int
|
| 657 |
+
if activeExpertsK > 0 {
|
| 658 |
+
activeK = activeExpertsK // user selected a value
|
| 659 |
+
} else {
|
| 660 |
+
// Auto-reduce for 397B on constrained devices
|
| 661 |
+
let is397B = model.path.lowercased().contains("397b") || model.name.lowercased().contains("397b")
|
| 662 |
+
let deviceRAM = ProcessInfo.processInfo.physicalMemory / (1024 * 1024 * 1024)
|
| 663 |
+
activeK = (is397B && deviceRAM <= 16) ? 4 : 0
|
| 664 |
+
}
|
| 665 |
+
|
| 666 |
+
// Decode ropeScaling tag: tens digit = mode, ones digit encodes factor
|
| 667 |
+
// 0 = off, 12 = linear 2x, 14 = linear 4x, 22 = NTK 2x, 24 = NTK 4x, 32 = YaRN 2x, 34 = YaRN 4x
|
| 668 |
+
let ropeMode = ropeScaling / 10
|
| 669 |
+
let ropeFactor: Float = ropeScaling == 0 ? 1.0 : Float(ropeScaling % 10)
|
| 670 |
+
|
| 671 |
+
Task {
|
| 672 |
+
do {
|
| 673 |
+
try await engine.loadModel(
|
| 674 |
+
at: model.path,
|
| 675 |
+
maxContext: maxContext,
|
| 676 |
+
thinkBudget: thinkingEnabled ? thinkBudget : -1,
|
| 677 |
+
useTiered: model.hasTiered,
|
| 678 |
+
activeExpertsK: activeK,
|
| 679 |
+
cacheIOSplit: cacheIOSplit,
|
| 680 |
+
cmdMerge: cmdMergeEnabled,
|
| 681 |
+
fusedAttention: fusedAttention,
|
| 682 |
+
expertPrefetch: expertPrefetch,
|
| 683 |
+
fusedExpert: fusedExpert,
|
| 684 |
+
fp16Accumulation: fp16Accumulation,
|
| 685 |
+
fp8KVCache: fp8KVCache,
|
| 686 |
+
slidingWindow: slidingWindow,
|
| 687 |
+
h2oBudget: h2oBudget,
|
| 688 |
+
ropeScalingMode: ropeMode,
|
| 689 |
+
ropeScaleFactor: ropeFactor,
|
| 690 |
+
prefillBatch: prefillBatch,
|
| 691 |
+
verbose: true
|
| 692 |
+
)
|
| 693 |
+
} catch {
|
| 694 |
+
// Error state is set by the engine
|
| 695 |
+
}
|
| 696 |
+
}
|
| 697 |
+
}
|
| 698 |
+
|
| 699 |
+
// MARK: - Custom URL Download
|
| 700 |
+
|
| 701 |
+
private func startCustomDownload() {
|
| 702 |
+
showCustomURLError = nil
|
| 703 |
+
isResolvingURL = true
|
| 704 |
+
var input = customRepoURL.trimmingCharacters(in: .whitespacesAndNewlines)
|
| 705 |
+
|
| 706 |
+
// Accept full URLs or repo IDs
|
| 707 |
+
// https://huggingface.co/user/repo → user/repo
|
| 708 |
+
if input.hasPrefix("https://huggingface.co/") {
|
| 709 |
+
input = String(input.dropFirst("https://huggingface.co/".count))
|
| 710 |
+
}
|
| 711 |
+
if input.hasPrefix("http://huggingface.co/") {
|
| 712 |
+
input = String(input.dropFirst("http://huggingface.co/".count))
|
| 713 |
+
}
|
| 714 |
+
// Remove trailing slashes and tree/main suffix
|
| 715 |
+
input = input.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
|
| 716 |
+
if input.hasSuffix("/tree/main") {
|
| 717 |
+
input = String(input.dropLast("/tree/main".count))
|
| 718 |
+
}
|
| 719 |
+
|
| 720 |
+
let parts = input.split(separator: "/")
|
| 721 |
+
guard parts.count >= 2 else {
|
| 722 |
+
showCustomURLError = "Invalid format. Use 'user/model' or a HuggingFace URL."
|
| 723 |
+
isResolvingURL = false
|
| 724 |
+
return
|
| 725 |
+
}
|
| 726 |
+
let repoId = "\(parts[0])/\(parts[1])"
|
| 727 |
+
let modelName = String(parts[1])
|
| 728 |
+
|
| 729 |
+
// Query HuggingFace API for file list
|
| 730 |
+
let apiURL = URL(string: "https://huggingface.co/api/models/\(repoId)")!
|
| 731 |
+
Task {
|
| 732 |
+
do {
|
| 733 |
+
let (data, response) = try await URLSession.shared.data(from: apiURL)
|
| 734 |
+
guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else {
|
| 735 |
+
await MainActor.run { showCustomURLError = "Repository not found: \(repoId)"; isResolvingURL = false }
|
| 736 |
+
return
|
| 737 |
+
}
|
| 738 |
+
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
|
| 739 |
+
let siblings = json["siblings"] as? [[String: Any]] else {
|
| 740 |
+
await MainActor.run { showCustomURLError = "Failed to parse repository metadata."; isResolvingURL = false }
|
| 741 |
+
return
|
| 742 |
+
}
|
| 743 |
+
|
| 744 |
+
var files: [RepoFile] = []
|
| 745 |
+
var totalSize: UInt64 = 0
|
| 746 |
+
var expertLayers = 0
|
| 747 |
+
|
| 748 |
+
for sibling in siblings {
|
| 749 |
+
guard let filename = sibling["rfilename"] as? String else { continue }
|
| 750 |
+
// Skip README, .gitattributes, etc.
|
| 751 |
+
let lower = filename.lowercased()
|
| 752 |
+
if lower.hasSuffix(".md") || lower.hasSuffix(".gitattributes") || lower.hasPrefix(".") { continue }
|
| 753 |
+
|
| 754 |
+
// Get file size from LFS info or default to 0
|
| 755 |
+
let size: UInt64
|
| 756 |
+
if let lfs = sibling["lfs"] as? [String: Any], let s = lfs["size"] as? UInt64 {
|
| 757 |
+
size = s
|
| 758 |
+
} else if let s = sibling["size"] as? UInt64 {
|
| 759 |
+
size = s
|
| 760 |
+
} else {
|
| 761 |
+
size = 0
|
| 762 |
+
}
|
| 763 |
+
|
| 764 |
+
files.append(RepoFile(filename: filename, sizeBytes: size))
|
| 765 |
+
totalSize += size
|
| 766 |
+
|
| 767 |
+
if filename.contains("packed_experts") && filename.hasSuffix(".bin") && filename.contains("layer_") {
|
| 768 |
+
expertLayers += 1
|
| 769 |
+
}
|
| 770 |
+
}
|
| 771 |
+
|
| 772 |
+
guard !files.isEmpty else {
|
| 773 |
+
await MainActor.run { showCustomURLError = "No files found in repository."; isResolvingURL = false }
|
| 774 |
+
return
|
| 775 |
+
}
|
| 776 |
+
|
| 777 |
+
// Check for required files
|
| 778 |
+
let hasConfig = files.contains { $0.filename == "config.json" }
|
| 779 |
+
let hasWeights = files.contains { $0.filename == "model_weights.bin" }
|
| 780 |
+
|| files.contains { $0.filename == "model_weights_0.bin" }
|
| 781 |
+
guard hasConfig && hasWeights else {
|
| 782 |
+
await MainActor.run {
|
| 783 |
+
showCustomURLError = "Not a Flash-MoE model (missing config.json or model_weights.bin)."
|
| 784 |
+
isResolvingURL = false
|
| 785 |
+
}
|
| 786 |
+
return
|
| 787 |
+
}
|
| 788 |
+
|
| 789 |
+
let entry = CatalogEntry(
|
| 790 |
+
id: "custom-\(modelName.lowercased())",
|
| 791 |
+
displayName: modelName,
|
| 792 |
+
repoId: repoId,
|
| 793 |
+
description: "Custom model from \(repoId)",
|
| 794 |
+
totalSizeBytes: totalSize,
|
| 795 |
+
quantization: "unknown",
|
| 796 |
+
expertLayers: expertLayers,
|
| 797 |
+
defaultK: 0,
|
| 798 |
+
recommendedK: 0,
|
| 799 |
+
minRAMGB: 0,
|
| 800 |
+
files: files
|
| 801 |
+
)
|
| 802 |
+
|
| 803 |
+
await MainActor.run {
|
| 804 |
+
customRepoURL = ""
|
| 805 |
+
showCustomURLError = nil
|
| 806 |
+
isResolvingURL = false
|
| 807 |
+
// Add to the download list — user taps the row to start
|
| 808 |
+
if !customEntries.contains(where: { $0.id == entry.id }) {
|
| 809 |
+
customEntries.append(entry)
|
| 810 |
+
}
|
| 811 |
+
}
|
| 812 |
+
} catch {
|
| 813 |
+
await MainActor.run { showCustomURLError = "Network error: \(error.localizedDescription)"; isResolvingURL = false }
|
| 814 |
+
}
|
| 815 |
+
}
|
| 816 |
+
}
|
| 817 |
+
|
| 818 |
+
// MARK: - Setting Label with Info Button
|
| 819 |
+
|
| 820 |
+
@ViewBuilder
|
| 821 |
+
private func settingLabel(_ title: String, key: String) -> some View {
|
| 822 |
+
HStack(spacing: 6) {
|
| 823 |
+
Button {
|
| 824 |
+
settingInfo = SettingInfo.all[key]
|
| 825 |
+
} label: {
|
| 826 |
+
Image(systemName: "info.circle")
|
| 827 |
+
.font(.caption)
|
| 828 |
+
.foregroundStyle(.blue)
|
| 829 |
+
}
|
| 830 |
+
.buttonStyle(.plain)
|
| 831 |
+
Text(title)
|
| 832 |
+
}
|
| 833 |
+
}
|
| 834 |
+
|
| 835 |
+
private func deleteModel(_ model: LocalModel) {
|
| 836 |
+
do {
|
| 837 |
+
try FileManager.default.removeItem(atPath: model.path)
|
| 838 |
+
print("[delete] Removed \(model.name) at \(model.path)")
|
| 839 |
+
scanForModels()
|
| 840 |
+
} catch {
|
| 841 |
+
print("ERROR: Failed to delete \(model.name): \(error)")
|
| 842 |
+
}
|
| 843 |
+
}
|
| 844 |
+
|
| 845 |
+
#if os(iOS)
|
| 846 |
+
// MARK: - File Import (iOS)
|
| 847 |
+
|
| 848 |
+
private func handleImportedFolder(_ url: URL) {
|
| 849 |
+
// Save a security-scoped bookmark so we can access this folder across launches
|
| 850 |
+
guard url.startAccessingSecurityScopedResource() else {
|
| 851 |
+
print("ERROR: Failed to access security-scoped resource")
|
| 852 |
+
return
|
| 853 |
+
}
|
| 854 |
+
|
| 855 |
+
do {
|
| 856 |
+
let bookmarkData = try url.bookmarkData(
|
| 857 |
+
options: .minimalBookmark,
|
| 858 |
+
includingResourceValuesForKeys: nil,
|
| 859 |
+
relativeTo: nil
|
| 860 |
+
)
|
| 861 |
+
// Save bookmark to UserDefaults
|
| 862 |
+
var bookmarks = UserDefaults.standard.array(forKey: "importedModelBookmarks") as? [Data] ?? []
|
| 863 |
+
bookmarks.append(bookmarkData)
|
| 864 |
+
UserDefaults.standard.set(bookmarks, forKey: "importedModelBookmarks")
|
| 865 |
+
|
| 866 |
+
print("[import] Bookmarked external model folder: \(url.path)")
|
| 867 |
+
scanForModels()
|
| 868 |
+
} catch {
|
| 869 |
+
print("ERROR: Failed to create bookmark: \(error)")
|
| 870 |
+
}
|
| 871 |
+
|
| 872 |
+
url.stopAccessingSecurityScopedResource()
|
| 873 |
+
}
|
| 874 |
+
|
| 875 |
+
private func moveImportedFolderToDocuments(_ url: URL) {
|
| 876 |
+
guard url.startAccessingSecurityScopedResource() else {
|
| 877 |
+
print("ERROR: Failed to access security-scoped resource for move")
|
| 878 |
+
return
|
| 879 |
+
}
|
| 880 |
+
|
| 881 |
+
let fm = FileManager.default
|
| 882 |
+
guard let docsDir = fm.urls(for: .documentDirectory, in: .userDomainMask).first else {
|
| 883 |
+
url.stopAccessingSecurityScopedResource()
|
| 884 |
+
return
|
| 885 |
+
}
|
| 886 |
+
|
| 887 |
+
let destURL = docsDir.appendingPathComponent(url.lastPathComponent)
|
| 888 |
+
|
| 889 |
+
importProgress = "Moving model to Documents..."
|
| 890 |
+
print("[import] Moving \(url.path) -> \(destURL.path)")
|
| 891 |
+
|
| 892 |
+
Task {
|
| 893 |
+
do {
|
| 894 |
+
// moveItem is instant if same filesystem, otherwise it copies
|
| 895 |
+
try fm.moveItem(at: url, to: destURL)
|
| 896 |
+
print("[import] Move succeeded")
|
| 897 |
+
} catch {
|
| 898 |
+
print("[import] Move failed: \(error). Trying copy...")
|
| 899 |
+
await MainActor.run { importProgress = "Copying model to Documents (this may take a while)..." }
|
| 900 |
+
do {
|
| 901 |
+
try fm.copyItem(at: url, to: destURL)
|
| 902 |
+
print("[import] Copy succeeded")
|
| 903 |
+
} catch {
|
| 904 |
+
print("ERROR: Copy also failed: \(error)")
|
| 905 |
+
}
|
| 906 |
+
}
|
| 907 |
+
|
| 908 |
+
url.stopAccessingSecurityScopedResource()
|
| 909 |
+
|
| 910 |
+
await MainActor.run {
|
| 911 |
+
importProgress = nil
|
| 912 |
+
scanForModels()
|
| 913 |
+
}
|
| 914 |
+
}
|
| 915 |
+
}
|
| 916 |
+
|
| 917 |
+
private func restoreBookmarks() {
|
| 918 |
+
guard let bookmarks = UserDefaults.standard.array(forKey: "importedModelBookmarks") as? [Data] else { return }
|
| 919 |
+
|
| 920 |
+
for bookmark in bookmarks {
|
| 921 |
+
var isStale = false
|
| 922 |
+
if let url = try? URL(resolvingBookmarkData: bookmark, bookmarkDataIsStale: &isStale) {
|
| 923 |
+
if !isStale {
|
| 924 |
+
_ = url.startAccessingSecurityScopedResource()
|
| 925 |
+
}
|
| 926 |
+
}
|
| 927 |
+
}
|
| 928 |
+
}
|
| 929 |
+
|
| 930 |
+
private func moveModelToExternal(model: LocalModel, destination: URL) {
|
| 931 |
+
Task {
|
| 932 |
+
let fm = FileManager.default
|
| 933 |
+
let destPath = destination.appendingPathComponent(URL(fileURLWithPath: model.path).lastPathComponent)
|
| 934 |
+
|
| 935 |
+
guard destination.startAccessingSecurityScopedResource() else {
|
| 936 |
+
print("ERROR: Cannot access destination")
|
| 937 |
+
return
|
| 938 |
+
}
|
| 939 |
+
defer { destination.stopAccessingSecurityScopedResource() }
|
| 940 |
+
|
| 941 |
+
do {
|
| 942 |
+
// Move (not copy) — instant on same filesystem
|
| 943 |
+
try fm.moveItem(at: URL(fileURLWithPath: model.path), to: destPath)
|
| 944 |
+
print("[export] Moved \(model.name) to \(destPath.path)")
|
| 945 |
+
await MainActor.run { scanForModels() }
|
| 946 |
+
} catch {
|
| 947 |
+
print("ERROR: Move failed: \(error). Trying copy instead...")
|
| 948 |
+
// If move fails (cross-volume), this would be slow for 300GB
|
| 949 |
+
// but at least it works
|
| 950 |
+
do {
|
| 951 |
+
try fm.copyItem(at: URL(fileURLWithPath: model.path), to: destPath)
|
| 952 |
+
try fm.removeItem(at: URL(fileURLWithPath: model.path))
|
| 953 |
+
print("[export] Copied + deleted \(model.name) to \(destPath.path)")
|
| 954 |
+
await MainActor.run { scanForModels() }
|
| 955 |
+
} catch {
|
| 956 |
+
print("ERROR: Copy also failed: \(error)")
|
| 957 |
+
}
|
| 958 |
+
}
|
| 959 |
+
}
|
| 960 |
+
}
|
| 961 |
+
#endif
|
| 962 |
+
}
|
| 963 |
+
|
| 964 |
+
// MARK: - Folder Import/Export Pickers (iOS)
|
| 965 |
+
|
| 966 |
+
#if os(iOS)
|
| 967 |
+
struct FolderImportPicker: UIViewControllerRepresentable {
|
| 968 |
+
let onPick: (URL) -> Void
|
| 969 |
+
|
| 970 |
+
func makeCoordinator() -> Coordinator { Coordinator(onPick: onPick) }
|
| 971 |
+
|
| 972 |
+
func makeUIViewController(context: Context) -> UIDocumentPickerViewController {
|
| 973 |
+
let picker = UIDocumentPickerViewController(forOpeningContentTypes: [.folder])
|
| 974 |
+
picker.allowsMultipleSelection = false
|
| 975 |
+
picker.delegate = context.coordinator
|
| 976 |
+
return picker
|
| 977 |
+
}
|
| 978 |
+
|
| 979 |
+
func updateUIViewController(_ uiViewController: UIDocumentPickerViewController, context: Context) {}
|
| 980 |
+
|
| 981 |
+
class Coordinator: NSObject, UIDocumentPickerDelegate {
|
| 982 |
+
let onPick: (URL) -> Void
|
| 983 |
+
init(onPick: @escaping (URL) -> Void) { self.onPick = onPick }
|
| 984 |
+
|
| 985 |
+
func documentPicker(_ controller: UIDocumentPickerViewController, didPickDocumentsAt urls: [URL]) {
|
| 986 |
+
guard let url = urls.first else { return }
|
| 987 |
+
onPick(url)
|
| 988 |
+
}
|
| 989 |
+
}
|
| 990 |
+
}
|
| 991 |
+
|
| 992 |
+
// MARK: - Folder Export Picker (pick destination to move model to)
|
| 993 |
+
|
| 994 |
+
struct FolderExportPicker: UIViewControllerRepresentable {
|
| 995 |
+
let sourceURL: URL
|
| 996 |
+
let onPick: (URL) -> Void
|
| 997 |
+
|
| 998 |
+
func makeCoordinator() -> Coordinator { Coordinator(onPick: onPick) }
|
| 999 |
+
|
| 1000 |
+
func makeUIViewController(context: Context) -> UIDocumentPickerViewController {
|
| 1001 |
+
// moveToService: shows full Files browser, user picks destination folder.
|
| 1002 |
+
// iOS moves the directory to the chosen location.
|
| 1003 |
+
let picker = UIDocumentPickerViewController(urls: [sourceURL], in: .moveToService)
|
| 1004 |
+
picker.delegate = context.coordinator
|
| 1005 |
+
return picker
|
| 1006 |
+
}
|
| 1007 |
+
|
| 1008 |
+
func updateUIViewController(_ uiViewController: UIDocumentPickerViewController, context: Context) {}
|
| 1009 |
+
|
| 1010 |
+
class Coordinator: NSObject, UIDocumentPickerDelegate {
|
| 1011 |
+
let onPick: (URL) -> Void
|
| 1012 |
+
init(onPick: @escaping (URL) -> Void) { self.onPick = onPick }
|
| 1013 |
+
|
| 1014 |
+
func documentPicker(_ controller: UIDocumentPickerViewController, didPickDocumentsAt urls: [URL]) {
|
| 1015 |
+
guard let url = urls.first else { return }
|
| 1016 |
+
onPick(url)
|
| 1017 |
+
}
|
| 1018 |
+
}
|
| 1019 |
+
}
|
| 1020 |
+
#endif
|
| 1021 |
+
|
| 1022 |
+
// MARK: - Model Row
|
| 1023 |
+
|
| 1024 |
+
struct ModelRow: View {
|
| 1025 |
+
let model: LocalModel
|
| 1026 |
+
let isLoading: Bool
|
| 1027 |
+
|
| 1028 |
+
var body: some View {
|
| 1029 |
+
HStack {
|
| 1030 |
+
VStack(alignment: .leading, spacing: 4) {
|
| 1031 |
+
Text(model.name)
|
| 1032 |
+
.font(.headline)
|
| 1033 |
+
|
| 1034 |
+
HStack(spacing: 8) {
|
| 1035 |
+
if model.hasTiered {
|
| 1036 |
+
QuantBadge(text: "Tiered", color: .green)
|
| 1037 |
+
} else if model.has4bit {
|
| 1038 |
+
QuantBadge(text: "4-bit", color: .blue)
|
| 1039 |
+
} else if model.has2bit {
|
| 1040 |
+
QuantBadge(text: "2-bit", color: .orange)
|
| 1041 |
+
}
|
| 1042 |
+
|
| 1043 |
+
Text(String(format: "%.1f GB", model.sizeGB))
|
| 1044 |
+
.font(.caption)
|
| 1045 |
+
.foregroundStyle(.secondary)
|
| 1046 |
+
}
|
| 1047 |
+
}
|
| 1048 |
+
|
| 1049 |
+
Spacer()
|
| 1050 |
+
|
| 1051 |
+
if isLoading {
|
| 1052 |
+
ProgressView()
|
| 1053 |
+
} else {
|
| 1054 |
+
Image(systemName: "chevron.right")
|
| 1055 |
+
.foregroundStyle(.secondary)
|
| 1056 |
+
}
|
| 1057 |
+
}
|
| 1058 |
+
.padding(.vertical, 4)
|
| 1059 |
+
.contentShape(Rectangle())
|
| 1060 |
+
}
|
| 1061 |
+
}
|
| 1062 |
+
|
| 1063 |
+
struct QuantBadge: View {
|
| 1064 |
+
let text: String
|
| 1065 |
+
let color: Color
|
| 1066 |
+
|
| 1067 |
+
var body: some View {
|
| 1068 |
+
Text(text)
|
| 1069 |
+
.font(.caption2.bold())
|
| 1070 |
+
.padding(.horizontal, 6)
|
| 1071 |
+
.padding(.vertical, 2)
|
| 1072 |
+
.background(color.opacity(0.15))
|
| 1073 |
+
.foregroundStyle(color)
|
| 1074 |
+
.clipShape(Capsule())
|
| 1075 |
+
}
|
| 1076 |
+
}
|
| 1077 |
+
|
| 1078 |
+
// MARK: - Model Scanner
|
| 1079 |
+
|
| 1080 |
+
enum ModelScanner {
|
| 1081 |
+
/// Scan common locations for Flash-MoE model directories
|
| 1082 |
+
static func scanLocalModels() async -> [LocalModel] {
|
| 1083 |
+
var models: [LocalModel] = []
|
| 1084 |
+
let fm = FileManager.default
|
| 1085 |
+
|
| 1086 |
+
// Scan app Documents directory
|
| 1087 |
+
if let docsDir = fm.urls(for: .documentDirectory, in: .userDomainMask).first {
|
| 1088 |
+
await scanDirectory(docsDir.path, into: &models)
|
| 1089 |
+
}
|
| 1090 |
+
|
| 1091 |
+
#if os(macOS)
|
| 1092 |
+
// Scan user-added model directories (macOS)
|
| 1093 |
+
if let macPaths = UserDefaults.standard.stringArray(forKey: "macModelPaths") {
|
| 1094 |
+
for path in macPaths {
|
| 1095 |
+
let fm2 = FileManager.default
|
| 1096 |
+
var isDir: ObjCBool = false
|
| 1097 |
+
guard fm2.fileExists(atPath: path, isDirectory: &isDir), isDir.boolValue else { continue }
|
| 1098 |
+
if FlashMoEEngine.validateModel(at: path) {
|
| 1099 |
+
let size = directorySize(at: path)
|
| 1100 |
+
let hasTiered = fm2.fileExists(atPath: (path as NSString).appendingPathComponent("packed_experts_tiered/layer_00.bin"))
|
| 1101 |
+
let has4bit = fm2.fileExists(atPath: (path as NSString).appendingPathComponent("packed_experts/layer_00.bin"))
|
| 1102 |
+
let has2bit = fm2.fileExists(atPath: (path as NSString).appendingPathComponent("packed_experts_2bit/layer_00.bin"))
|
| 1103 |
+
models.append(LocalModel(
|
| 1104 |
+
name: URL(fileURLWithPath: path).lastPathComponent,
|
| 1105 |
+
path: path,
|
| 1106 |
+
sizeBytes: size,
|
| 1107 |
+
hasTiered: hasTiered,
|
| 1108 |
+
has4bit: has4bit,
|
| 1109 |
+
has2bit: has2bit
|
| 1110 |
+
))
|
| 1111 |
+
} else {
|
| 1112 |
+
await scanDirectory(path, into: &models)
|
| 1113 |
+
}
|
| 1114 |
+
}
|
| 1115 |
+
}
|
| 1116 |
+
#endif
|
| 1117 |
+
|
| 1118 |
+
// Scan bookmarked external folders (imported via Files picker on iOS)
|
| 1119 |
+
if let bookmarks = UserDefaults.standard.array(forKey: "importedModelBookmarks") as? [Data] {
|
| 1120 |
+
for bookmark in bookmarks {
|
| 1121 |
+
var isStale = false
|
| 1122 |
+
if let url = try? URL(resolvingBookmarkData: bookmark, bookmarkDataIsStale: &isStale),
|
| 1123 |
+
!isStale {
|
| 1124 |
+
let accessed = url.startAccessingSecurityScopedResource()
|
| 1125 |
+
// Check if the folder itself is a model
|
| 1126 |
+
if FlashMoEEngine.validateModel(at: url.path) {
|
| 1127 |
+
let size = directorySize(at: url.path)
|
| 1128 |
+
let hasTiered = fm.fileExists(atPath: url.appendingPathComponent("packed_experts_tiered/layer_00.bin").path)
|
| 1129 |
+
let has4bit = fm.fileExists(atPath: url.appendingPathComponent("packed_experts/layer_00.bin").path)
|
| 1130 |
+
let has2bit = fm.fileExists(atPath: url.appendingPathComponent("packed_experts_2bit/layer_00.bin").path)
|
| 1131 |
+
models.append(LocalModel(
|
| 1132 |
+
name: "📁 " + url.lastPathComponent,
|
| 1133 |
+
path: url.path,
|
| 1134 |
+
sizeBytes: size,
|
| 1135 |
+
hasTiered: hasTiered,
|
| 1136 |
+
has4bit: has4bit,
|
| 1137 |
+
has2bit: has2bit
|
| 1138 |
+
))
|
| 1139 |
+
} else {
|
| 1140 |
+
// Scan subdirectories
|
| 1141 |
+
await scanDirectory(url.path, into: &models)
|
| 1142 |
+
}
|
| 1143 |
+
if accessed { url.stopAccessingSecurityScopedResource() }
|
| 1144 |
+
}
|
| 1145 |
+
}
|
| 1146 |
+
}
|
| 1147 |
+
|
| 1148 |
+
return models.sorted { $0.name < $1.name }
|
| 1149 |
+
}
|
| 1150 |
+
|
| 1151 |
+
private static func scanDirectory(_ path: String, into models: inout [LocalModel]) async {
|
| 1152 |
+
let fm = FileManager.default
|
| 1153 |
+
|
| 1154 |
+
guard let entries = try? fm.contentsOfDirectory(atPath: path) else { return }
|
| 1155 |
+
|
| 1156 |
+
for entry in entries {
|
| 1157 |
+
let fullPath = (path as NSString).appendingPathComponent(entry)
|
| 1158 |
+
var isDir: ObjCBool = false
|
| 1159 |
+
guard fm.fileExists(atPath: fullPath, isDirectory: &isDir), isDir.boolValue else { continue }
|
| 1160 |
+
|
| 1161 |
+
// Check if it's a valid model
|
| 1162 |
+
if FlashMoEEngine.validateModel(at: fullPath) {
|
| 1163 |
+
// Protect model files from iOS storage optimization / purging
|
| 1164 |
+
excludeFromBackup(URL(fileURLWithPath: fullPath))
|
| 1165 |
+
let size = directorySize(at: fullPath)
|
| 1166 |
+
let hasTiered = fm.fileExists(atPath: (fullPath as NSString).appendingPathComponent("packed_experts_tiered/layer_00.bin"))
|
| 1167 |
+
let has4bit = fm.fileExists(atPath: (fullPath as NSString).appendingPathComponent("packed_experts/layer_00.bin"))
|
| 1168 |
+
let has2bit = fm.fileExists(atPath: (fullPath as NSString).appendingPathComponent("packed_experts_2bit/layer_00.bin"))
|
| 1169 |
+
|
| 1170 |
+
models.append(LocalModel(
|
| 1171 |
+
name: entry,
|
| 1172 |
+
path: fullPath,
|
| 1173 |
+
sizeBytes: size,
|
| 1174 |
+
hasTiered: hasTiered,
|
| 1175 |
+
has4bit: has4bit,
|
| 1176 |
+
has2bit: has2bit
|
| 1177 |
+
))
|
| 1178 |
+
}
|
| 1179 |
+
}
|
| 1180 |
+
}
|
| 1181 |
+
|
| 1182 |
+
private static func directorySize(at path: String) -> UInt64 {
|
| 1183 |
+
let fm = FileManager.default
|
| 1184 |
+
guard let enumerator = fm.enumerator(atPath: path) else { return 0 }
|
| 1185 |
+
var total: UInt64 = 0
|
| 1186 |
+
while let file = enumerator.nextObject() as? String {
|
| 1187 |
+
let fullPath = (path as NSString).appendingPathComponent(file)
|
| 1188 |
+
if let attrs = try? fm.attributesOfItem(atPath: fullPath),
|
| 1189 |
+
let size = attrs[.size] as? UInt64 {
|
| 1190 |
+
total += size
|
| 1191 |
+
}
|
| 1192 |
+
}
|
| 1193 |
+
return total
|
| 1194 |
+
}
|
| 1195 |
+
|
| 1196 |
+
/// Mark a directory (and its contents) as excluded from iCloud backup and
|
| 1197 |
+
/// iOS storage optimization, preventing the system from purging model files.
|
| 1198 |
+
private static func excludeFromBackup(_ url: URL) {
|
| 1199 |
+
var url = url
|
| 1200 |
+
var values = URLResourceValues()
|
| 1201 |
+
values.isExcludedFromBackup = true
|
| 1202 |
+
try? url.setResourceValues(values)
|
| 1203 |
+
|
| 1204 |
+
// Also mark all files inside
|
| 1205 |
+
let fm = FileManager.default
|
| 1206 |
+
guard let enumerator = fm.enumerator(at: url, includingPropertiesForKeys: nil) else { return }
|
| 1207 |
+
while let fileURL = enumerator.nextObject() as? URL {
|
| 1208 |
+
var fileURL = fileURL
|
| 1209 |
+
try? fileURL.setResourceValues(values)
|
| 1210 |
+
}
|
| 1211 |
+
}
|
| 1212 |
+
}
|
FlashMoE-iOS/Views/ProfilerView.swift
ADDED
|
@@ -0,0 +1,274 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/*
|
| 2 |
+
* ProfilerView.swift — Lightweight resource profiler overlay
|
| 3 |
+
*
|
| 4 |
+
* Displays real-time system metrics during inference:
|
| 5 |
+
* memory (RSS + available), thermal state, CPU usage,
|
| 6 |
+
* and engine stats (tok/s, TTFT).
|
| 7 |
+
*
|
| 8 |
+
* All APIs are public (mach_task_info, os_proc_available_memory,
|
| 9 |
+
* ProcessInfo.thermalState) — no entitlements needed.
|
| 10 |
+
*/
|
| 11 |
+
|
| 12 |
+
import SwiftUI
|
| 13 |
+
import Darwin.Mach
|
| 14 |
+
|
| 15 |
+
// MARK: - System Metrics Sampler
|
| 16 |
+
|
| 17 |
+
@Observable
|
| 18 |
+
final class SystemMetrics: @unchecked Sendable {
|
| 19 |
+
private(set) var residentMemoryMB: Double = 0
|
| 20 |
+
private(set) var availableMemoryMB: Double = 0
|
| 21 |
+
private(set) var cpuUsagePercent: Double = 0
|
| 22 |
+
private(set) var thermalState: ProcessInfo.ThermalState = .nominal
|
| 23 |
+
|
| 24 |
+
private var prevCPUTime: Double = 0
|
| 25 |
+
private var prevSampleTime: CFAbsoluteTime = 0
|
| 26 |
+
|
| 27 |
+
func sample() {
|
| 28 |
+
residentMemoryMB = Self.getResidentMemory()
|
| 29 |
+
availableMemoryMB = Self.getAvailableMemory()
|
| 30 |
+
cpuUsagePercent = sampleCPU()
|
| 31 |
+
thermalState = ProcessInfo.processInfo.thermalState
|
| 32 |
+
}
|
| 33 |
+
|
| 34 |
+
// MARK: - Memory
|
| 35 |
+
|
| 36 |
+
private static func getResidentMemory() -> Double {
|
| 37 |
+
var info = mach_task_basic_info()
|
| 38 |
+
var count = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size / MemoryLayout<natural_t>.size)
|
| 39 |
+
let result = withUnsafeMutablePointer(to: &info) {
|
| 40 |
+
$0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
|
| 41 |
+
task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &count)
|
| 42 |
+
}
|
| 43 |
+
}
|
| 44 |
+
guard result == KERN_SUCCESS else { return 0 }
|
| 45 |
+
return Double(info.resident_size) / (1024 * 1024)
|
| 46 |
+
}
|
| 47 |
+
|
| 48 |
+
private static func getAvailableMemory() -> Double {
|
| 49 |
+
#if os(iOS)
|
| 50 |
+
// iOS: use os_proc_available_memory()
|
| 51 |
+
return Double(os_proc_available_memory()) / (1024 * 1024)
|
| 52 |
+
#elseif os(macOS)
|
| 53 |
+
// macOS: estimate available memory using host_statistics64 and page counts
|
| 54 |
+
var vmStats = vm_statistics64()
|
| 55 |
+
var count = mach_msg_type_number_t(MemoryLayout<vm_statistics64_data_t>.size / MemoryLayout<integer_t>.size)
|
| 56 |
+
var size: vm_size_t = 0
|
| 57 |
+
let host = mach_host_self()
|
| 58 |
+
|
| 59 |
+
// Get page size
|
| 60 |
+
let kerrPage = host_page_size(host, &size)
|
| 61 |
+
guard kerrPage == KERN_SUCCESS else { return 0 }
|
| 62 |
+
|
| 63 |
+
// Fetch VM statistics
|
| 64 |
+
let result: kern_return_t = withUnsafeMutablePointer(to: &vmStats) { ptr in
|
| 65 |
+
ptr.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { intPtr in
|
| 66 |
+
host_statistics64(host, HOST_VM_INFO64, intPtr, &count)
|
| 67 |
+
}
|
| 68 |
+
}
|
| 69 |
+
guard result == KERN_SUCCESS else { return 0 }
|
| 70 |
+
|
| 71 |
+
// Consider free + inactive pages as "available"
|
| 72 |
+
let freePages = UInt64(vmStats.free_count)
|
| 73 |
+
let inactivePages = UInt64(vmStats.inactive_count)
|
| 74 |
+
let speculativePages = UInt64(vmStats.speculative_count)
|
| 75 |
+
|
| 76 |
+
let availableBytes = (freePages + inactivePages + speculativePages) * UInt64(size)
|
| 77 |
+
return Double(availableBytes) / (1024 * 1024)
|
| 78 |
+
#else
|
| 79 |
+
// Other platforms: not available — return 0
|
| 80 |
+
return 0
|
| 81 |
+
#endif
|
| 82 |
+
}
|
| 83 |
+
|
| 84 |
+
// MARK: - CPU
|
| 85 |
+
|
| 86 |
+
private func sampleCPU() -> Double {
|
| 87 |
+
let now = CFAbsoluteTimeGetCurrent()
|
| 88 |
+
let totalCPU = Self.getThreadCPUTime()
|
| 89 |
+
|
| 90 |
+
defer {
|
| 91 |
+
prevCPUTime = totalCPU
|
| 92 |
+
prevSampleTime = now
|
| 93 |
+
}
|
| 94 |
+
|
| 95 |
+
guard prevSampleTime > 0 else { return 0 }
|
| 96 |
+
let elapsed = now - prevSampleTime
|
| 97 |
+
guard elapsed > 0 else { return 0 }
|
| 98 |
+
|
| 99 |
+
let cpuDelta = totalCPU - prevCPUTime
|
| 100 |
+
// Normalize to percentage (cpuDelta is in seconds of CPU time)
|
| 101 |
+
return min((cpuDelta / elapsed) * 100.0, 999.0)
|
| 102 |
+
}
|
| 103 |
+
|
| 104 |
+
private static func getThreadCPUTime() -> Double {
|
| 105 |
+
var threadList: thread_act_array_t?
|
| 106 |
+
var threadCount: mach_msg_type_number_t = 0
|
| 107 |
+
let result = task_threads(mach_task_self_, &threadList, &threadCount)
|
| 108 |
+
guard result == KERN_SUCCESS, let threads = threadList else { return 0 }
|
| 109 |
+
defer {
|
| 110 |
+
vm_deallocate(mach_task_self_, vm_address_t(bitPattern: threads), vm_size_t(Int(threadCount) * MemoryLayout<thread_t>.size))
|
| 111 |
+
}
|
| 112 |
+
|
| 113 |
+
var total: Double = 0
|
| 114 |
+
for i in 0..<Int(threadCount) {
|
| 115 |
+
var info = thread_basic_info()
|
| 116 |
+
var infoCount = mach_msg_type_number_t(MemoryLayout<thread_basic_info_data_t>.size / MemoryLayout<natural_t>.size)
|
| 117 |
+
let kr = withUnsafeMutablePointer(to: &info) {
|
| 118 |
+
$0.withMemoryRebound(to: integer_t.self, capacity: Int(infoCount)) {
|
| 119 |
+
thread_info(threads[i], thread_flavor_t(THREAD_BASIC_INFO), $0, &infoCount)
|
| 120 |
+
}
|
| 121 |
+
}
|
| 122 |
+
if kr == KERN_SUCCESS {
|
| 123 |
+
total += Double(info.user_time.seconds) + Double(info.user_time.microseconds) / 1_000_000
|
| 124 |
+
total += Double(info.system_time.seconds) + Double(info.system_time.microseconds) / 1_000_000
|
| 125 |
+
}
|
| 126 |
+
}
|
| 127 |
+
return total
|
| 128 |
+
}
|
| 129 |
+
}
|
| 130 |
+
|
| 131 |
+
// MARK: - Profiler View
|
| 132 |
+
|
| 133 |
+
struct ProfilerView: View {
|
| 134 |
+
let engine: FlashMoEEngine
|
| 135 |
+
@State private var metrics = SystemMetrics()
|
| 136 |
+
@State private var timer: Timer?
|
| 137 |
+
|
| 138 |
+
var body: some View {
|
| 139 |
+
VStack(spacing: 0) {
|
| 140 |
+
// Header
|
| 141 |
+
HStack {
|
| 142 |
+
Image(systemName: "gauge.with.dots.needle.50percent")
|
| 143 |
+
.foregroundStyle(.orange)
|
| 144 |
+
Text("Profiler")
|
| 145 |
+
.font(.caption.bold())
|
| 146 |
+
Spacer()
|
| 147 |
+
thermalBadge
|
| 148 |
+
}
|
| 149 |
+
.padding(.horizontal, 12)
|
| 150 |
+
.padding(.top, 8)
|
| 151 |
+
.padding(.bottom, 4)
|
| 152 |
+
|
| 153 |
+
Divider().opacity(0.3)
|
| 154 |
+
|
| 155 |
+
// Metrics grid
|
| 156 |
+
LazyVGrid(columns: [
|
| 157 |
+
GridItem(.flexible()),
|
| 158 |
+
GridItem(.flexible()),
|
| 159 |
+
], spacing: 6) {
|
| 160 |
+
metricCell(
|
| 161 |
+
icon: "memorychip",
|
| 162 |
+
label: "RSS",
|
| 163 |
+
value: String(format: "%.0f MB", metrics.residentMemoryMB)
|
| 164 |
+
)
|
| 165 |
+
metricCell(
|
| 166 |
+
icon: "memorychip.fill",
|
| 167 |
+
label: "Free",
|
| 168 |
+
value: String(format: "%.0f MB", metrics.availableMemoryMB)
|
| 169 |
+
)
|
| 170 |
+
metricCell(
|
| 171 |
+
icon: "cpu",
|
| 172 |
+
label: "CPU",
|
| 173 |
+
value: String(format: "%.0f%%", metrics.cpuUsagePercent)
|
| 174 |
+
)
|
| 175 |
+
metricCell(
|
| 176 |
+
icon: "speedometer",
|
| 177 |
+
label: "tok/s",
|
| 178 |
+
value: String(format: "%.1f", engine.tokensPerSecond)
|
| 179 |
+
)
|
| 180 |
+
metricCell(
|
| 181 |
+
icon: "number",
|
| 182 |
+
label: "Tokens",
|
| 183 |
+
value: "\(engine.tokensGenerated)"
|
| 184 |
+
)
|
| 185 |
+
metricCell(
|
| 186 |
+
icon: "timer",
|
| 187 |
+
label: "TTFT",
|
| 188 |
+
value: engine.timeToFirstToken > 0
|
| 189 |
+
? (engine.timeToFirstToken > 500_000
|
| 190 |
+
? String(format: "%.1f min", engine.timeToFirstToken / 60_000)
|
| 191 |
+
: String(format: "%.0f ms", engine.timeToFirstToken))
|
| 192 |
+
: "--"
|
| 193 |
+
)
|
| 194 |
+
}
|
| 195 |
+
.padding(.horizontal, 12)
|
| 196 |
+
.padding(.vertical, 8)
|
| 197 |
+
}
|
| 198 |
+
.background(.ultraThinMaterial)
|
| 199 |
+
.clipShape(RoundedRectangle(cornerRadius: 12))
|
| 200 |
+
.shadow(color: .black.opacity(0.15), radius: 8, y: 2)
|
| 201 |
+
.padding(.horizontal)
|
| 202 |
+
.onAppear { startSampling() }
|
| 203 |
+
.onDisappear { stopSampling() }
|
| 204 |
+
}
|
| 205 |
+
|
| 206 |
+
// MARK: - Subviews
|
| 207 |
+
|
| 208 |
+
private func metricCell(icon: String, label: String, value: String) -> some View {
|
| 209 |
+
HStack(spacing: 6) {
|
| 210 |
+
Image(systemName: icon)
|
| 211 |
+
.font(.caption2)
|
| 212 |
+
.foregroundStyle(.secondary)
|
| 213 |
+
.frame(width: 14)
|
| 214 |
+
VStack(alignment: .leading, spacing: 0) {
|
| 215 |
+
Text(label)
|
| 216 |
+
.font(.system(size: 9))
|
| 217 |
+
.foregroundStyle(.tertiary)
|
| 218 |
+
Text(value)
|
| 219 |
+
.font(.caption.monospacedDigit().bold())
|
| 220 |
+
}
|
| 221 |
+
Spacer()
|
| 222 |
+
}
|
| 223 |
+
}
|
| 224 |
+
|
| 225 |
+
private var thermalBadge: some View {
|
| 226 |
+
HStack(spacing: 3) {
|
| 227 |
+
Circle()
|
| 228 |
+
.fill(thermalColor)
|
| 229 |
+
.frame(width: 6, height: 6)
|
| 230 |
+
Text(thermalLabel)
|
| 231 |
+
.font(.system(size: 9).bold())
|
| 232 |
+
.foregroundStyle(thermalColor)
|
| 233 |
+
}
|
| 234 |
+
.padding(.horizontal, 6)
|
| 235 |
+
.padding(.vertical, 2)
|
| 236 |
+
.background(thermalColor.opacity(0.12))
|
| 237 |
+
.clipShape(Capsule())
|
| 238 |
+
}
|
| 239 |
+
|
| 240 |
+
private var thermalColor: Color {
|
| 241 |
+
switch metrics.thermalState {
|
| 242 |
+
case .nominal: return .green
|
| 243 |
+
case .fair: return .yellow
|
| 244 |
+
case .serious: return .orange
|
| 245 |
+
case .critical: return .red
|
| 246 |
+
@unknown default: return .gray
|
| 247 |
+
}
|
| 248 |
+
}
|
| 249 |
+
|
| 250 |
+
private var thermalLabel: String {
|
| 251 |
+
switch metrics.thermalState {
|
| 252 |
+
case .nominal: return "Cool"
|
| 253 |
+
case .fair: return "Warm"
|
| 254 |
+
case .serious: return "Hot"
|
| 255 |
+
case .critical: return "Critical"
|
| 256 |
+
@unknown default: return "?"
|
| 257 |
+
}
|
| 258 |
+
}
|
| 259 |
+
|
| 260 |
+
// MARK: - Sampling
|
| 261 |
+
|
| 262 |
+
private func startSampling() {
|
| 263 |
+
metrics.sample() // initial
|
| 264 |
+
timer = Timer.scheduledTimer(withTimeInterval: 1.0, repeats: true) { _ in
|
| 265 |
+
metrics.sample()
|
| 266 |
+
}
|
| 267 |
+
}
|
| 268 |
+
|
| 269 |
+
private func stopSampling() {
|
| 270 |
+
timer?.invalidate()
|
| 271 |
+
timer = nil
|
| 272 |
+
}
|
| 273 |
+
}
|
| 274 |
+
|
FlashMoE-iOS/project.yml
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: FlashMoE
|
| 2 |
+
options:
|
| 3 |
+
bundleIdPrefix: com.flashmoe
|
| 4 |
+
deploymentTarget:
|
| 5 |
+
iOS: "18.0"
|
| 6 |
+
xcodeVersion: "26.0"
|
| 7 |
+
generateEmptyDirectories: true
|
| 8 |
+
|
| 9 |
+
settings:
|
| 10 |
+
base:
|
| 11 |
+
SWIFT_VERSION: "6.0"
|
| 12 |
+
IPHONEOS_DEPLOYMENT_TARGET: "18.0"
|
| 13 |
+
TARGETED_DEVICE_FAMILY: "1" # iPhone only
|
| 14 |
+
INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace: YES
|
| 15 |
+
INFOPLIST_KEY_UIFileSharingEnabled: YES
|
| 16 |
+
|
| 17 |
+
targets:
|
| 18 |
+
FlashMoE:
|
| 19 |
+
type: application
|
| 20 |
+
platform: iOS
|
| 21 |
+
sources:
|
| 22 |
+
- path: App
|
| 23 |
+
type: group
|
| 24 |
+
- path: Views
|
| 25 |
+
type: group
|
| 26 |
+
- path: ViewModels
|
| 27 |
+
type: group
|
| 28 |
+
- path: Bridge
|
| 29 |
+
type: group
|
| 30 |
+
- path: FlashMoEEngine
|
| 31 |
+
type: group
|
| 32 |
+
- path: Resources
|
| 33 |
+
type: group
|
| 34 |
+
# Metal shaders — copied as resource (compiled at runtime via newLibraryWithSource)
|
| 35 |
+
- path: ../metal_infer/shaders.metal
|
| 36 |
+
type: file
|
| 37 |
+
buildPhase: resources
|
| 38 |
+
# Tokenizer header — needed by infer.m
|
| 39 |
+
- path: ../metal_infer/tokenizer.h
|
| 40 |
+
type: file
|
| 41 |
+
buildPhase: headers
|
| 42 |
+
settings:
|
| 43 |
+
base:
|
| 44 |
+
PRODUCT_BUNDLE_IDENTIFIER: com.flashmoe.ios
|
| 45 |
+
PRODUCT_NAME: FlashMoE
|
| 46 |
+
SWIFT_OBJC_BRIDGING_HEADER: Bridge/FlashMoE-Bridging-Header.h
|
| 47 |
+
# Obj-C / C compilation flags for infer.m unity build
|
| 48 |
+
OTHER_CFLAGS: "-DCHAT_MODE=1 -DACCELERATE_NEW_LAPACK"
|
| 49 |
+
# Header search paths
|
| 50 |
+
HEADER_SEARCH_PATHS:
|
| 51 |
+
- $(SRCROOT)/FlashMoEEngine
|
| 52 |
+
- $(SRCROOT)/../metal_infer
|
| 53 |
+
# Needed for the engine
|
| 54 |
+
CLANG_ENABLE_OBJC_ARC: YES
|
| 55 |
+
GCC_OPTIMIZATION_LEVEL: "2"
|
| 56 |
+
GCC_WARN_INHIBIT_ALL_WARNINGS: NO
|
| 57 |
+
# Memory / entitlements
|
| 58 |
+
INFOPLIST_KEY_UILaunchScreen_Generation: YES
|
| 59 |
+
ASSETCATALOG_COMPILER_APPICON_NAME: AppIcon
|
| 60 |
+
# Info.plist keys
|
| 61 |
+
INFOPLIST_KEY_CFBundleDisplayName: "Flash-MoE"
|
| 62 |
+
INFOPLIST_KEY_UIApplicationSceneManifest_Generation: YES
|
| 63 |
+
INFOPLIST_KEY_UISupportedInterfaceOrientations: "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"
|
| 64 |
+
# Linker flags for compression library
|
| 65 |
+
OTHER_LDFLAGS: "-lcompression"
|
| 66 |
+
dependencies:
|
| 67 |
+
- sdk: Metal.framework
|
| 68 |
+
- sdk: Foundation.framework
|
| 69 |
+
- sdk: Accelerate.framework
|
| 70 |
+
preBuildScripts: []
|
| 71 |
+
postBuildScripts: []
|
| 72 |
+
entitlements:
|
| 73 |
+
path: Resources/FlashMoE.entitlements
|
| 74 |
+
properties:
|
| 75 |
+
com.apple.security.app-sandbox: true
|
README.md
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
CLAUDE.md
|
app.py
ADDED
|
@@ -0,0 +1,164 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import tempfile
|
| 3 |
+
import subprocess
|
| 4 |
+
|
| 5 |
+
os.environ["HF_HUB_CACHE"] = "cache"
|
| 6 |
+
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
|
| 7 |
+
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
|
| 8 |
+
import gradio as gr
|
| 9 |
+
|
| 10 |
+
from huggingface_hub import HfApi
|
| 11 |
+
from huggingface_hub import whoami
|
| 12 |
+
from huggingface_hub import ModelCard
|
| 13 |
+
from huggingface_hub import scan_cache_dir
|
| 14 |
+
from huggingface_hub import snapshot_download
|
| 15 |
+
from huggingface_hub import logging
|
| 16 |
+
|
| 17 |
+
from gradio_huggingfacehub_search import HuggingfaceHubSearch
|
| 18 |
+
from apscheduler.schedulers.background import BackgroundScheduler
|
| 19 |
+
|
| 20 |
+
from textwrap import dedent
|
| 21 |
+
|
| 22 |
+
HF_TOKEN = os.environ.get("HF_TOKEN")
|
| 23 |
+
|
| 24 |
+
def clear_hf_cache_space():
|
| 25 |
+
scan = scan_cache_dir()
|
| 26 |
+
to_delete = []
|
| 27 |
+
for repo in scan.repos:
|
| 28 |
+
if repo.repo_type == "model":
|
| 29 |
+
to_delete.extend([rev.commit_hash for rev in repo.revisions])
|
| 30 |
+
scan.delete_revisions(*to_delete).execute()
|
| 31 |
+
print("Cache has been cleared")
|
| 32 |
+
|
| 33 |
+
def upload_to_hub(model_path, upload_repo, hf_path, oauth_token):
|
| 34 |
+
card = ModelCard(dedent(
|
| 35 |
+
f"""
|
| 36 |
+
# {upload_repo}
|
| 37 |
+
|
| 38 |
+
The Model [{upload_repo}](https://huggingface.co/{upload_repo}) was prepared for Flash-MoE inference from [{hf_path}](https://huggingface.co/{hf_path}).
|
| 39 |
+
|
| 40 |
+
## Use with Flash-MoE
|
| 41 |
+
|
| 42 |
+
```bash
|
| 43 |
+
# Clone the engine
|
| 44 |
+
git clone https://github.com/alexintosh/flash-moe.git
|
| 45 |
+
cd flash-moe/metal_infer
|
| 46 |
+
make
|
| 47 |
+
|
| 48 |
+
# Run inference
|
| 49 |
+
./infer --model {upload_repo} --prompt "Explain quantum computing" --tokens 100
|
| 50 |
+
```
|
| 51 |
+
"""
|
| 52 |
+
))
|
| 53 |
+
card.save(os.path.join(model_path, "README.md"))
|
| 54 |
+
|
| 55 |
+
logging.set_verbosity_info()
|
| 56 |
+
|
| 57 |
+
api = HfApi(token=oauth_token.token)
|
| 58 |
+
api.create_repo(repo_id=upload_repo, exist_ok=True)
|
| 59 |
+
|
| 60 |
+
print(f"Uploading files to: {upload_repo}")
|
| 61 |
+
|
| 62 |
+
# We ignore the original safetensors to save massive amounts of space,
|
| 63 |
+
# as Flash-MoE inference relies entirely on the packed bins and extracted weights.
|
| 64 |
+
api.upload_folder(
|
| 65 |
+
folder_path=model_path,
|
| 66 |
+
repo_id=upload_repo,
|
| 67 |
+
ignore_patterns=["*.safetensors", "*.safetensors.index.json"],
|
| 68 |
+
commit_message="Add Flash-MoE packed experts and extracted weights"
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
print(f"Upload successful, go to https://huggingface.co/{upload_repo} for details.")
|
| 72 |
+
|
| 73 |
+
def process_model(model_id, oauth_token: gr.OAuthToken | None):
|
| 74 |
+
if oauth_token is None or oauth_token.token is None:
|
| 75 |
+
raise ValueError("You must be logged in to use Flash-MoE Packager")
|
| 76 |
+
|
| 77 |
+
model_name = model_id.split('/')[-1]
|
| 78 |
+
username = whoami(oauth_token.token)["name"]
|
| 79 |
+
upload_repo = f"{username}/{model_name}-FlashMoE"
|
| 80 |
+
|
| 81 |
+
try:
|
| 82 |
+
print(f"Downloading model {model_id}...")
|
| 83 |
+
cache_dir = snapshot_download(model_id, token=oauth_token.token)
|
| 84 |
+
|
| 85 |
+
with tempfile.TemporaryDirectory(dir=".") as tmpdir:
|
| 86 |
+
model_path = os.path.join(tmpdir, "model")
|
| 87 |
+
os.makedirs(model_path, exist_ok=True)
|
| 88 |
+
|
| 89 |
+
# Symlink everything from the read-only HF cache into our writable tmp model dir
|
| 90 |
+
# This allows the flash-moe python scripts to safely write `packed_experts/` directly inside
|
| 91 |
+
for item in os.listdir(cache_dir):
|
| 92 |
+
src = os.path.join(cache_dir, item)
|
| 93 |
+
dst = os.path.join(model_path, item)
|
| 94 |
+
os.symlink(src, dst)
|
| 95 |
+
|
| 96 |
+
index_file = os.path.join(tmpdir, "expert_index.json")
|
| 97 |
+
|
| 98 |
+
# 1. Build expert index
|
| 99 |
+
print("1. Building expert index...")
|
| 100 |
+
subprocess.run(
|
| 101 |
+
["python", "build_expert_index.py", "--model", model_path, "--output", index_file],
|
| 102 |
+
check=True
|
| 103 |
+
)
|
| 104 |
+
|
| 105 |
+
# 2. Pack expert weights into per-layer binary files
|
| 106 |
+
print("2. Repacking experts...")
|
| 107 |
+
subprocess.run(
|
| 108 |
+
["python", "repack_experts.py", "--index", index_file],
|
| 109 |
+
check=True
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
# 3. Extract non-expert weights into a single mmap-friendly binary
|
| 113 |
+
print("3. Extracting non-expert weights...")
|
| 114 |
+
subprocess.run(
|
| 115 |
+
["python", "metal_infer/extract_weights.py", "--model", model_path, "--output", model_path],
|
| 116 |
+
check=True
|
| 117 |
+
)
|
| 118 |
+
|
| 119 |
+
print("Uploading to Hugging Face Hub...")
|
| 120 |
+
upload_to_hub(model_path, upload_repo, model_id, oauth_token)
|
| 121 |
+
|
| 122 |
+
return f'🚀 Success! Find your Flash-MoE ready repo <a href="https://hf.co/{upload_repo}" target="_blank" style="text-decoration:underline">here</a>'
|
| 123 |
+
|
| 124 |
+
except Exception as e:
|
| 125 |
+
return f"❌ Error: {e}"
|
| 126 |
+
finally:
|
| 127 |
+
clear_hf_cache_space()
|
| 128 |
+
print("Folder cleaned up successfully!")
|
| 129 |
+
|
| 130 |
+
css="""/* Custom CSS to allow scrolling */
|
| 131 |
+
.gradio-container {overflow-y: auto;}
|
| 132 |
+
"""
|
| 133 |
+
|
| 134 |
+
# Create Gradio interface
|
| 135 |
+
with gr.Blocks(css=css) as demo:
|
| 136 |
+
gr.Markdown("# ⚡ Flash-MoE Auto-Packager")
|
| 137 |
+
gr.Markdown("You must be logged in to package and push models.")
|
| 138 |
+
gr.LoginButton(min_width=250)
|
| 139 |
+
|
| 140 |
+
model_id = HuggingfaceHubSearch(
|
| 141 |
+
label="MLX Model ID",
|
| 142 |
+
placeholder="e.g. mlx-community/Qwen3.5-35B-A3B-4bit",
|
| 143 |
+
search_type="model",
|
| 144 |
+
)
|
| 145 |
+
|
| 146 |
+
iface = gr.Interface(
|
| 147 |
+
fn=process_model,
|
| 148 |
+
inputs=[model_id],
|
| 149 |
+
outputs=[gr.Markdown(label="Status")],
|
| 150 |
+
title="Package models for blazingly fast MoE inference! ⚡",
|
| 151 |
+
description="This Space downloads an MLX-quantized MoE model, extracts the safetensor indices, repacks the experts for Apple Silicon SSD streaming, and uploads the optimized weights to a new repo under your namespace.",
|
| 152 |
+
api_name=False,
|
| 153 |
+
allow_flagging="never"
|
| 154 |
+
)
|
| 155 |
+
|
| 156 |
+
def restart_space():
|
| 157 |
+
HfApi().restart_space(repo_id=os.environ.get("SPACE_ID", "alexintosh/flash-moe-packager"), token=HF_TOKEN, factory_reboot=True)
|
| 158 |
+
|
| 159 |
+
scheduler = BackgroundScheduler()
|
| 160 |
+
scheduler.add_job(restart_space, "interval", seconds=21600)
|
| 161 |
+
scheduler.start()
|
| 162 |
+
|
| 163 |
+
# Launch the interface
|
| 164 |
+
demo.queue(default_concurrency_limit=1, max_size=5).launch(debug=True, show_api=False)
|
autoresearch/baseline.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
5.62
|
autoresearch/benchmark.sh
ADDED
|
@@ -0,0 +1,121 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# benchmark.sh — Fixed harness for Flash-MoE autoresearch.
|
| 3 |
+
# DO NOT MODIFY. This is the ground truth measurement.
|
| 4 |
+
#
|
| 5 |
+
# Usage: bash autoresearch/benchmark.sh
|
| 6 |
+
# Requires: FLASH_MOE_MODEL env var set to model path
|
| 7 |
+
#
|
| 8 |
+
# Exit codes: 0=all pass, 1=build fail, 2=quality fail, 3=runtime error
|
| 9 |
+
|
| 10 |
+
set -euo pipefail
|
| 11 |
+
|
| 12 |
+
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
| 13 |
+
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
|
| 14 |
+
METAL_DIR="$PROJECT_DIR/metal_infer"
|
| 15 |
+
|
| 16 |
+
# --- Config ---
|
| 17 |
+
BENCH_TOKENS=200
|
| 18 |
+
BENCH_PROMPT="Explain quantum computing in simple terms"
|
| 19 |
+
MATH_PROMPT="What is 2+2? Answer with just the number."
|
| 20 |
+
JSON_PROMPT='You are a helpful assistant. Return a JSON object with keys "name" and "value" for the capital of France. Output ONLY valid JSON, nothing else.'
|
| 21 |
+
WARMUP_TOKENS=5
|
| 22 |
+
|
| 23 |
+
# --- Validate environment ---
|
| 24 |
+
if [ -z "${FLASH_MOE_MODEL:-}" ]; then
|
| 25 |
+
echo "ERROR: FLASH_MOE_MODEL not set" >&2
|
| 26 |
+
exit 1
|
| 27 |
+
fi
|
| 28 |
+
if [ ! -f "$FLASH_MOE_MODEL/config.json" ]; then
|
| 29 |
+
echo "ERROR: No config.json in $FLASH_MOE_MODEL" >&2
|
| 30 |
+
exit 1
|
| 31 |
+
fi
|
| 32 |
+
|
| 33 |
+
# --- Step 1: BUILD ---
|
| 34 |
+
echo "=== BUILD ===" >&2
|
| 35 |
+
cd "$METAL_DIR"
|
| 36 |
+
|
| 37 |
+
# Touch infer.m to ensure shaders.metal changes trigger rebuild
|
| 38 |
+
# (shaders are compiled at runtime from source embedded in infer.m's build)
|
| 39 |
+
touch infer.m
|
| 40 |
+
if ! make infer 2>&1 | tail -5 >&2; then
|
| 41 |
+
echo "BENCH_RESULT tok_s=0.00 math=SKIP json=SKIP status=BUILD_FAIL"
|
| 42 |
+
exit 1
|
| 43 |
+
fi
|
| 44 |
+
echo "Build OK" >&2
|
| 45 |
+
|
| 46 |
+
# --- Step 2: WARMUP (prime page cache) ---
|
| 47 |
+
echo "=== WARMUP ===" >&2
|
| 48 |
+
./infer --model "$FLASH_MOE_MODEL" --prompt "Hi" --tokens "$WARMUP_TOKENS" > /dev/null 2>&1 || true
|
| 49 |
+
|
| 50 |
+
# --- Step 3: PERFORMANCE BENCHMARK ---
|
| 51 |
+
echo "=== BENCHMARK ($BENCH_TOKENS tokens) ===" >&2
|
| 52 |
+
BENCH_OUT=$(mktemp)
|
| 53 |
+
BENCH_ERR=$(mktemp)
|
| 54 |
+
MATH_OUT=$(mktemp)
|
| 55 |
+
MATH_ERR=$(mktemp)
|
| 56 |
+
JSON_OUT=$(mktemp)
|
| 57 |
+
JSON_ERR=$(mktemp)
|
| 58 |
+
trap "rm -f $BENCH_OUT $BENCH_ERR $MATH_OUT $MATH_ERR $JSON_OUT $JSON_ERR 2>/dev/null" EXIT
|
| 59 |
+
|
| 60 |
+
if ! ./infer --model "$FLASH_MOE_MODEL" \
|
| 61 |
+
--prompt "$BENCH_PROMPT" \
|
| 62 |
+
--tokens "$BENCH_TOKENS" \
|
| 63 |
+
> "$BENCH_OUT" 2> "$BENCH_ERR"; then
|
| 64 |
+
echo "BENCH_RESULT tok_s=0.00 math=SKIP json=SKIP status=RUNTIME_ERROR"
|
| 65 |
+
cat "$BENCH_ERR" >&2
|
| 66 |
+
exit 3
|
| 67 |
+
fi
|
| 68 |
+
|
| 69 |
+
# Parse tok/s from "Generation: X.X s (Y.YY tok/s)"
|
| 70 |
+
TOK_S=$(grep "Generation:" "$BENCH_OUT" | grep -oE '[0-9]+\.[0-9]+ tok/s' | grep -oE '[0-9]+\.[0-9]+')
|
| 71 |
+
if [ -z "$TOK_S" ]; then
|
| 72 |
+
echo "ERROR: Could not parse tok/s from output" >&2
|
| 73 |
+
cat "$BENCH_OUT" >&2
|
| 74 |
+
echo "BENCH_RESULT tok_s=0.00 math=SKIP json=SKIP status=PARSE_ERROR"
|
| 75 |
+
exit 3
|
| 76 |
+
fi
|
| 77 |
+
echo "Performance: $TOK_S tok/s" >&2
|
| 78 |
+
|
| 79 |
+
# --- Step 4: QUALITY GATE — MATH ---
|
| 80 |
+
echo "=== QUALITY: MATH ===" >&2
|
| 81 |
+
MATH_PASS="FAIL"
|
| 82 |
+
if ./infer --model "$FLASH_MOE_MODEL" \
|
| 83 |
+
--prompt "$MATH_PROMPT" \
|
| 84 |
+
--tokens 20 \
|
| 85 |
+
--think-budget 256 \
|
| 86 |
+
> "$MATH_OUT" 2> "$MATH_ERR"; then
|
| 87 |
+
# Extract text between "--- Output ---" and "--- Statistics ---"
|
| 88 |
+
MATH_TEXT=$(sed -n '/--- Output ---/,/--- Statistics ---/{/--- Output ---/d;/--- Statistics ---/d;p;}' "$MATH_OUT")
|
| 89 |
+
if echo "$MATH_TEXT" | grep -q "4"; then
|
| 90 |
+
MATH_PASS="PASS"
|
| 91 |
+
fi
|
| 92 |
+
fi
|
| 93 |
+
echo "Math gate: $MATH_PASS" >&2
|
| 94 |
+
|
| 95 |
+
# --- Step 5: QUALITY GATE — JSON ---
|
| 96 |
+
echo "=== QUALITY: JSON ===" >&2
|
| 97 |
+
JSON_PASS="FAIL"
|
| 98 |
+
if ./infer --model "$FLASH_MOE_MODEL" \
|
| 99 |
+
--prompt "$JSON_PROMPT" \
|
| 100 |
+
--tokens 50 \
|
| 101 |
+
--think-budget 256 \
|
| 102 |
+
> "$JSON_OUT" 2> "$JSON_ERR"; then
|
| 103 |
+
JSON_TEXT=$(sed -n '/--- Output ---/,/--- Statistics ---/{/--- Output ---/d;/--- Statistics ---/d;p;}' "$JSON_OUT")
|
| 104 |
+
# Check for valid JSON structure: has braces and a key
|
| 105 |
+
if echo "$JSON_TEXT" | grep -q '{' && echo "$JSON_TEXT" | grep -q '}' && echo "$JSON_TEXT" | grep -q '"name"'; then
|
| 106 |
+
JSON_PASS="PASS"
|
| 107 |
+
fi
|
| 108 |
+
fi
|
| 109 |
+
echo "JSON gate: $JSON_PASS" >&2
|
| 110 |
+
|
| 111 |
+
# --- Step 6: STRUCTURED OUTPUT ---
|
| 112 |
+
if [ "$MATH_PASS" = "PASS" ] && [ "$JSON_PASS" = "PASS" ]; then
|
| 113 |
+
STATUS="OK"
|
| 114 |
+
EXIT_CODE=0
|
| 115 |
+
else
|
| 116 |
+
STATUS="QUALITY_FAIL"
|
| 117 |
+
EXIT_CODE=2
|
| 118 |
+
fi
|
| 119 |
+
|
| 120 |
+
echo "BENCH_RESULT tok_s=$TOK_S math=$MATH_PASS json=$JSON_PASS status=$STATUS"
|
| 121 |
+
exit $EXIT_CODE
|
autoresearch/experiments.tsv
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
id timestamp commit tok_s math json status notes
|
| 2 |
+
000 2026-03-21T01:45:41 0c2827f 5.62 PASS PASS baseline Unmodified baseline
|
| 3 |
+
001 2026-03-21T01:47:25 33bf261 5.63 PASS PASS keep SIMD reduction in rms_norm_qk + gated_rms_norm (replace serial thread-0 loop)
|
| 4 |
+
002 2026-03-21T01:48:54 a768791 5.97 PASS PASS keep FMA optimization in 2-bit dequant kernel — +6% from baseline
|
| 5 |
+
003 2026-03-21T01:50:57 e6c21c2 5.86 PASS PASS discard Single compute encoder for batch matvec — regression, GPU serializes
|
| 6 |
+
004 2026-03-21T01:52:45 ecfab24 5.37 PASS PASS discard FMA in matvec_fast — regression, no shared x cache means extra device reads
|
| 7 |
+
005 2026-03-21T01:55:04 4836dd8 5.76 PASS PASS discard SwiGLU vec4 — regression, poor occupancy at small dim
|
| 8 |
+
006 2026-03-21T01:57:37 39ff380 5.36 PASS PASS discard v3 kernel with 8192 shared mem — regression, occupancy drop from 32KB threadgroup mem
|
| 9 |
+
007 2026-03-21T01:59:47 b0a3e42 5.84 PASS PASS discard matvec_fast tg_size 64->128 — regression
|
| 10 |
+
008 2026-03-21T02:02:32 6c280cd 6.30 PASS PASS keep Half-precision x_shared in v3 kernel — 8KB shared mem, +12% from baseline
|
| 11 |
+
009 2026-03-21T02:04:02 1282be7 6.48 PASS PASS keep Half-precision x_shared in 2-bit kernel — +15.3% total from baseline
|
autoresearch/findings.md
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Autoresearch v2 Findings
|
| 2 |
+
|
| 3 |
+
Structured results from paper-guided experiments. Read this BEFORE starting new experiments to avoid repeating dead ends.
|
| 4 |
+
|
| 5 |
+
## Historical Context (from v1 and manual optimization)
|
| 6 |
+
|
| 7 |
+
Before this log, 58+ experiments were run. Key learnings:
|
| 8 |
+
|
| 9 |
+
- **FMA dequant kernel**: +12% from rearranging `(nibble * scale + bias) * x` → `fma(nibble, scale*x, bias*x)`. KEPT.
|
| 10 |
+
- **Trust OS page cache**: Every custom cache was slower. Metal LRU (-38%), malloc (-20%), LZ4 (-13%).
|
| 11 |
+
- **Deferred CMD3**: GPU/CPU overlap for expert forward pass. KEPT.
|
| 12 |
+
- **BLAS delta-net**: cblas_sscal/sgemv/sger for 64-head recurrence. +64% on attention. KEPT.
|
| 13 |
+
- **F_NOCACHE for 2-bit**: +3% from avoiding page cache thrash with smaller working set. 2-bit only.
|
| 14 |
+
- **GPU fused attention (RoPE)**: +2% for full-attention layers. KEPT.
|
| 15 |
+
- **C BPE tokenizer**: 180ms vs 3500ms startup. 20x improvement. KEPT.
|
| 16 |
+
- **Half-precision x_shared in 2-bit**: 6.48 tok/s (best from v1 experiments). KEPT.
|
| 17 |
+
|
| 18 |
+
### Failed approaches (DO NOT RETRY):
|
| 19 |
+
- LZ4 expert compression: -13% (decompression overhead > savings)
|
| 20 |
+
- F_RDADVISE prefetch: net 0% (unified memory: SSD DMA slows GPU -73%)
|
| 21 |
+
- Temporal expert prediction: -18% (25% hit rate, bandwidth waste)
|
| 22 |
+
- MLP routing predictor: 31% accuracy (worse than temporal)
|
| 23 |
+
- GPU LUT dequant kernel: -2% (indirect register access serializes)
|
| 24 |
+
- GPU private buffer compression: -20% (blit cost 4x7MB > savings)
|
| 25 |
+
- Spin-poll GPU wait: -23% (CPU thermal competes with GPU)
|
| 26 |
+
- Expert file clustering: 0% (NVMe ignores scatter at 7MB granularity)
|
| 27 |
+
- dispatch_io: -70% (dispatch_data management overhead)
|
| 28 |
+
- mmap expert files: -5x (per-page fault overhead on cold data)
|
| 29 |
+
- Speculative early routing: -38% (cache pollution + overhead)
|
| 30 |
+
- MTP speculative decoding: break-even (MoE I/O scales per-token)
|
| 31 |
+
|
| 32 |
+
---
|
| 33 |
+
|
| 34 |
+
_New experiments start below. Each entry follows the structured format._
|
| 35 |
+
|
autoresearch/prepare.sh
ADDED
|
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# prepare.sh — One-time setup for Flash-MoE autoresearch.
|
| 3 |
+
# Validates environment, builds, runs baseline benchmark.
|
| 4 |
+
#
|
| 5 |
+
# Usage: bash autoresearch/prepare.sh
|
| 6 |
+
|
| 7 |
+
set -euo pipefail
|
| 8 |
+
|
| 9 |
+
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
| 10 |
+
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
|
| 11 |
+
|
| 12 |
+
echo "========================================"
|
| 13 |
+
echo " Flash-MoE Autoresearch — Setup"
|
| 14 |
+
echo "========================================"
|
| 15 |
+
echo ""
|
| 16 |
+
|
| 17 |
+
# --- Check model ---
|
| 18 |
+
if [ -z "${FLASH_MOE_MODEL:-}" ]; then
|
| 19 |
+
echo "ERROR: Set FLASH_MOE_MODEL to your model path."
|
| 20 |
+
echo " export FLASH_MOE_MODEL=~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit"
|
| 21 |
+
exit 1
|
| 22 |
+
fi
|
| 23 |
+
|
| 24 |
+
if [ ! -f "$FLASH_MOE_MODEL/config.json" ]; then
|
| 25 |
+
echo "ERROR: No config.json in $FLASH_MOE_MODEL"
|
| 26 |
+
exit 1
|
| 27 |
+
fi
|
| 28 |
+
|
| 29 |
+
echo "Model: $FLASH_MOE_MODEL"
|
| 30 |
+
echo ""
|
| 31 |
+
|
| 32 |
+
# --- Check source files ---
|
| 33 |
+
for f in metal_infer/infer.m metal_infer/shaders.metal metal_infer/Makefile; do
|
| 34 |
+
if [ ! -f "$PROJECT_DIR/$f" ]; then
|
| 35 |
+
echo "ERROR: Missing $f"
|
| 36 |
+
exit 1
|
| 37 |
+
fi
|
| 38 |
+
done
|
| 39 |
+
echo "Source files: OK"
|
| 40 |
+
|
| 41 |
+
# --- Build ---
|
| 42 |
+
echo ""
|
| 43 |
+
echo "Building..."
|
| 44 |
+
cd "$PROJECT_DIR/metal_infer"
|
| 45 |
+
make clean && make infer
|
| 46 |
+
echo "Build: OK"
|
| 47 |
+
echo ""
|
| 48 |
+
|
| 49 |
+
# --- Check for existing baseline ---
|
| 50 |
+
BASELINE_FILE="$SCRIPT_DIR/baseline.txt"
|
| 51 |
+
TSV_FILE="$SCRIPT_DIR/experiments.tsv"
|
| 52 |
+
|
| 53 |
+
if [ -f "$BASELINE_FILE" ]; then
|
| 54 |
+
EXISTING=$(cat "$BASELINE_FILE")
|
| 55 |
+
echo "Existing baseline found: $EXISTING tok/s"
|
| 56 |
+
echo -n "Overwrite? [y/N] "
|
| 57 |
+
read -r REPLY
|
| 58 |
+
if [ "$REPLY" != "y" ] && [ "$REPLY" != "Y" ]; then
|
| 59 |
+
echo "Keeping existing baseline."
|
| 60 |
+
exit 0
|
| 61 |
+
fi
|
| 62 |
+
fi
|
| 63 |
+
|
| 64 |
+
# --- Run baseline ---
|
| 65 |
+
echo "Running baseline benchmark..."
|
| 66 |
+
echo ""
|
| 67 |
+
|
| 68 |
+
RESULT=$(bash "$SCRIPT_DIR/benchmark.sh")
|
| 69 |
+
echo "$RESULT"
|
| 70 |
+
|
| 71 |
+
# Parse tok/s
|
| 72 |
+
TOK_S=$(echo "$RESULT" | grep "BENCH_RESULT" | grep -oE 'tok_s=[0-9.]+' | cut -d= -f2)
|
| 73 |
+
MATH=$(echo "$RESULT" | grep "BENCH_RESULT" | grep -oE 'math=[A-Z]+' | cut -d= -f2)
|
| 74 |
+
JSON=$(echo "$RESULT" | grep "BENCH_RESULT" | grep -oE 'json=[A-Z]+' | cut -d= -f2)
|
| 75 |
+
STATUS=$(echo "$RESULT" | grep "BENCH_RESULT" | grep -oE 'status=[A-Z_]+' | cut -d= -f2)
|
| 76 |
+
|
| 77 |
+
if [ "$STATUS" != "OK" ]; then
|
| 78 |
+
echo ""
|
| 79 |
+
echo "ERROR: Baseline failed with status=$STATUS"
|
| 80 |
+
echo "Fix the issue before starting autoresearch."
|
| 81 |
+
exit 1
|
| 82 |
+
fi
|
| 83 |
+
|
| 84 |
+
# Save baseline
|
| 85 |
+
echo "$TOK_S" > "$BASELINE_FILE"
|
| 86 |
+
|
| 87 |
+
# Record current commit
|
| 88 |
+
COMMIT=$(git -C "$PROJECT_DIR" rev-parse --short HEAD)
|
| 89 |
+
TIMESTAMP=$(date -u +%Y-%m-%dT%H:%M:%S)
|
| 90 |
+
|
| 91 |
+
# Initialize TSV
|
| 92 |
+
echo -e "id\ttimestamp\tcommit\ttok_s\tmath\tjson\tstatus\tnotes" > "$TSV_FILE"
|
| 93 |
+
echo -e "000\t$TIMESTAMP\t$COMMIT\t$TOK_S\t$MATH\t$JSON\tbaseline\tUnmodified baseline" >> "$TSV_FILE"
|
| 94 |
+
|
| 95 |
+
echo ""
|
| 96 |
+
echo "========================================"
|
| 97 |
+
echo " Baseline: $TOK_S tok/s"
|
| 98 |
+
echo " Math: $MATH JSON: $JSON"
|
| 99 |
+
echo " Saved to: autoresearch/baseline.txt"
|
| 100 |
+
echo " TSV: autoresearch/experiments.tsv"
|
| 101 |
+
echo "========================================"
|
| 102 |
+
echo ""
|
| 103 |
+
echo "Ready. Start the agent with program.md."
|
autoresearch/program.md
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Flash-MoE Autoresearch
|
| 2 |
+
|
| 3 |
+
Autonomous optimization of MoE inference throughput on Apple Silicon.
|
| 4 |
+
|
| 5 |
+
## Objective
|
| 6 |
+
|
| 7 |
+
**Maximize tok/s** on the fixed benchmark while maintaining output quality. The benchmark generates 200 tokens with a fixed prompt. Your metric is `tok_s` — higher is better.
|
| 8 |
+
|
| 9 |
+
## Setup
|
| 10 |
+
|
| 11 |
+
To set up a new experiment run, work with the user to:
|
| 12 |
+
|
| 13 |
+
1. **Agree on a run tag**: propose a tag based on today's date (e.g. `mar21`). The branch `autoresearch/<tag>` must not already exist.
|
| 14 |
+
2. **Create the branch**: `git checkout -b autoresearch/<tag>` from the current branch.
|
| 15 |
+
3. **Read the in-scope files** for full context:
|
| 16 |
+
- `CLAUDE.md` — architecture overview, what worked, what failed.
|
| 17 |
+
- `metal_infer/infer.m` — the inference engine (~7500 lines, Objective-C). This is your primary target.
|
| 18 |
+
- `metal_infer/shaders.metal` — Metal compute kernels (~1300 lines). Secondary target.
|
| 19 |
+
- `autoresearch/benchmark.sh` — the fixed benchmark harness. Do not modify.
|
| 20 |
+
4. **Verify model**: Confirm `FLASH_MOE_MODEL` is set and `autoresearch/baseline.txt` exists. If not, tell the human to run `bash autoresearch/prepare.sh`.
|
| 21 |
+
5. **Initialize experiments.tsv**: If starting fresh, create it with the header row plus baseline from `baseline.txt`.
|
| 22 |
+
6. **Confirm and go**: Confirm setup looks good, then start experimenting.
|
| 23 |
+
|
| 24 |
+
## Files You May Modify
|
| 25 |
+
|
| 26 |
+
- `metal_infer/infer.m` — the inference engine. Everything is fair game: pipeline scheduling, buffer management, dispatch patterns, kernel launch parameters, memory layout, data flow.
|
| 27 |
+
- `metal_infer/shaders.metal` — Metal compute kernels. Fair game: threadgroup sizes, tiling strategy, SIMD utilization, register pressure, kernel fusion, shared memory patterns.
|
| 28 |
+
|
| 29 |
+
## Files You Must NOT Modify
|
| 30 |
+
|
| 31 |
+
- `autoresearch/benchmark.sh` — the measurement harness is sacred
|
| 32 |
+
- `autoresearch/prepare.sh`
|
| 33 |
+
- `metal_infer/Makefile`
|
| 34 |
+
- `metal_infer/chat.m`, `main.m`, `tokenizer.h`
|
| 35 |
+
- Any Python files (`*.py`)
|
| 36 |
+
- `CLAUDE.md`
|
| 37 |
+
|
| 38 |
+
## The Experiment Loop
|
| 39 |
+
|
| 40 |
+
Each experiment: modify → build → benchmark → decide → repeat. Each iteration takes ~1-2 minutes.
|
| 41 |
+
|
| 42 |
+
**LOOP FOREVER:**
|
| 43 |
+
|
| 44 |
+
1. **Propose** a single, focused optimization hypothesis. Write it down.
|
| 45 |
+
2. **Implement** the change in `infer.m` and/or `shaders.metal`.
|
| 46 |
+
3. **Commit**: `git add metal_infer/infer.m metal_infer/shaders.metal && git commit -m "<description>"`
|
| 47 |
+
4. **Benchmark**: `bash autoresearch/benchmark.sh 2>/dev/null` (stderr has progress, stdout has the BENCH_RESULT line)
|
| 48 |
+
- To see progress: `bash autoresearch/benchmark.sh` (without redirect)
|
| 49 |
+
5. **Parse** the BENCH_RESULT line: `tok_s=X.XX math=PASS/FAIL json=PASS/FAIL status=OK/...`
|
| 50 |
+
6. **If crashed or build failed**: `tail -20` the build output, attempt a quick fix. If unfixable after 2 tries, revert and log as `crash`.
|
| 51 |
+
7. **Record** in `autoresearch/experiments.tsv`
|
| 52 |
+
8. **Decide**:
|
| 53 |
+
- **KEEP** if `tok_s >= previous_best * 0.995` AND both quality gates pass. Advance the branch.
|
| 54 |
+
- **DISCARD** if `tok_s < previous_best * 0.995` OR any quality gate fails. Revert: `git reset --hard HEAD~1`
|
| 55 |
+
9. **Repeat**
|
| 56 |
+
|
| 57 |
+
## Architecture Constraints — READ CAREFULLY
|
| 58 |
+
|
| 59 |
+
These are hard-won lessons from 58+ prior experiments. Violating them will waste your time.
|
| 60 |
+
|
| 61 |
+
### The Unified Memory Constraint
|
| 62 |
+
|
| 63 |
+
On Apple Silicon, **SSD DMA and GPU compute share the same memory controller**. They cannot be profitably overlapped. The GPU's dequant kernels are bandwidth-saturated at ~418 GiB/s. Even small background SSD DMA causes disproportionate GPU latency spikes through memory controller arbitration. The serial pipeline (GPU → SSD → GPU) is hardware-optimal.
|
| 64 |
+
|
| 65 |
+
**Do NOT attempt:**
|
| 66 |
+
- Overlapping SSD reads with GPU compute
|
| 67 |
+
- dispatch_io (70% slower due to dispatch_data overhead)
|
| 68 |
+
- F_RDADVISE / speculative prefetch (net 0% — bandwidth waste)
|
| 69 |
+
- mmap for expert files (5x slower — per-page fault overhead)
|
| 70 |
+
|
| 71 |
+
### The Caching Constraint
|
| 72 |
+
|
| 73 |
+
**Trust the OS page cache.** Every custom expert cache we tried was slower:
|
| 74 |
+
- Metal LRU cache: -38% (steals GPU memory)
|
| 75 |
+
- malloc cache: -20% (steals from page cache)
|
| 76 |
+
- LZ4 compressed cache: -13% (decompression overhead)
|
| 77 |
+
- Speculative early routing: -38% (cache pollution)
|
| 78 |
+
|
| 79 |
+
The OS page cache (~35GB) does LRU better than we can. Do not add caching layers.
|
| 80 |
+
|
| 81 |
+
### GPU Constraints
|
| 82 |
+
|
| 83 |
+
- Metal GPU is **memory-bandwidth-bound** at ~418 GiB/s. The dequant kernels are already saturated. Pure ALU tricks (LUT, bit shifts) don't help because the bottleneck is reading data.
|
| 84 |
+
- **No spin-polling.** CPU thermal throttling competes with GPU on unified architecture. Use proper `waitUntilCompleted` or command buffer completion handlers.
|
| 85 |
+
- Expert files are ~6.75MB each. NVMe doesn't care about scatter at this granularity.
|
| 86 |
+
|
| 87 |
+
### What Already Works Well
|
| 88 |
+
|
| 89 |
+
Don't reinvent these — they're already optimized:
|
| 90 |
+
- **FMA dequant kernel**: `fma(nibble, scale*x, bias*x)` — pre-computes `scale*x` and `bias*x`
|
| 91 |
+
- **Deferred CMD3**: expert forward pass submitted without waiting, GPU/CPU overlap
|
| 92 |
+
- **BLAS delta-net**: Accelerate framework for the 64-head state recurrence
|
| 93 |
+
- **GCD parallel pread**: dispatch groups for K=4 expert I/O
|
| 94 |
+
- **GPU fused attention**: RoPE + QK norm fused
|
| 95 |
+
- **Fused moe_combine_residual**: single kernel for combine + residual + sigmoid gate
|
| 96 |
+
|
| 97 |
+
## Promising Areas to Explore
|
| 98 |
+
|
| 99 |
+
Ideas that haven't been tried or were partially explored:
|
| 100 |
+
|
| 101 |
+
1. **Shader occupancy tuning** — threadgroup sizes, simdgroup counts, register pressure in the dequant kernels. The current defaults may not be optimal for all matrix dimensions.
|
| 102 |
+
|
| 103 |
+
2. **Command buffer batching** — reducing Metal encode overhead by batching more operations into fewer command buffers. Currently 3 command buffers per layer.
|
| 104 |
+
|
| 105 |
+
3. **Pipeline stage rebalancing** — the bottleneck is expert I/O at 2.41ms per layer. Can any GPU work be restructured to reduce total wall time?
|
| 106 |
+
|
| 107 |
+
4. **Reduced precision in non-critical paths** — bf16/f16 for intermediate buffers where full f32 isn't needed. Apple GPU has native f16 ALU at 2x throughput.
|
| 108 |
+
|
| 109 |
+
5. **Kernel fusion** — combining small sequential kernels (e.g., norm + projection) to reduce dispatch overhead and intermediate buffer traffic.
|
| 110 |
+
|
| 111 |
+
6. **Attention layer optimization** — the full-attention layers use batched GPU attention. Room for improvement in tiling or memory access patterns.
|
| 112 |
+
|
| 113 |
+
7. **Thread coarsening in dequant kernels** — processing more output elements per thread to amortize the shared memory load of the input vector.
|
| 114 |
+
|
| 115 |
+
8. **SIMD shuffle patterns** — using simd_shuffle instead of threadgroup memory for reductions in the dequant kernels.
|
| 116 |
+
|
| 117 |
+
9. **Encode parallelism** — can command buffer encoding happen on a background thread while the previous buffer executes?
|
| 118 |
+
|
| 119 |
+
10. **Quantization-aware optimizations** — the 4-bit nibble extraction uses shifts and masks; explore whether packed byte operations are faster.
|
| 120 |
+
|
| 121 |
+
## Recording Results
|
| 122 |
+
|
| 123 |
+
Append to `autoresearch/experiments.tsv` (tab-separated, NOT comma-separated):
|
| 124 |
+
|
| 125 |
+
```
|
| 126 |
+
id timestamp commit tok_s math json status notes
|
| 127 |
+
```
|
| 128 |
+
|
| 129 |
+
Example:
|
| 130 |
+
```
|
| 131 |
+
001 2026-03-21T02:15:00 abc1234 4.42 PASS PASS keep Increased threadgroup size in matvec_v3 from 256 to 512
|
| 132 |
+
002 2026-03-21T02:17:30 def5678 4.35 PASS PASS discard Tried simd_shuffle for reduction — register pressure
|
| 133 |
+
003 2026-03-21T02:19:45 ghi9012 0.00 SKIP SKIP crash Fused norm+proj kernel — indexing bug
|
| 134 |
+
```
|
| 135 |
+
|
| 136 |
+
**Do NOT commit experiments.tsv** — leave it untracked by git.
|
| 137 |
+
|
| 138 |
+
## Decision Rules
|
| 139 |
+
|
| 140 |
+
- **Keep** if: `tok_s >= previous_best * 0.995` AND both quality gates PASS (0.5% noise margin to avoid discarding neutral-but-correct simplifications)
|
| 141 |
+
- **Discard** if: `tok_s < previous_best * 0.995` OR any quality gate FAIL
|
| 142 |
+
- If **3 consecutive discards** with similar approaches: move to a different optimization area
|
| 143 |
+
- If **build fails**: quick fix attempt (2 tries max), then revert and move on
|
| 144 |
+
- After every **keep**: update your mental baseline to the new tok_s
|
| 145 |
+
- Always **revert before starting a new experiment** (clean slate from last keep)
|
| 146 |
+
|
| 147 |
+
## Simplicity Criterion
|
| 148 |
+
|
| 149 |
+
All else being equal, simpler is better. A 0.5% improvement that adds 50 lines of complexity? Borderline. A 0.5% improvement from *removing* code? Definitely keep. An improvement of ~0 but simpler code? Keep. The goal is a lean, fast engine — not a pile of micro-optimizations.
|
| 150 |
+
|
| 151 |
+
## Safety
|
| 152 |
+
|
| 153 |
+
- This is a primary dev machine with 48GB unified RAM
|
| 154 |
+
- Do not allocate more than ~200MB of new Metal buffers
|
| 155 |
+
- Do not create files larger than 10MB
|
| 156 |
+
- Do not modify the build system or add dependencies
|
| 157 |
+
- If build fails, revert immediately
|
| 158 |
+
- Do not touch files outside `metal_infer/infer.m` and `metal_infer/shaders.metal`
|
| 159 |
+
|
| 160 |
+
## NEVER STOP
|
| 161 |
+
|
| 162 |
+
Once the loop begins, do NOT pause to ask the human for permission. Do NOT ask "should I keep going?" or "is this a good stopping point?". The human may be asleep. You run autonomously until manually interrupted. If you run out of ideas, think harder — re-read the source for new angles, try combining near-misses, try more radical changes. Each experiment takes ~1-2 minutes, so you can run ~30-60 per hour. The loop runs until the human stops you.
|
autoresearch/program_v2.md
ADDED
|
@@ -0,0 +1,212 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Flash-MoE Autoresearch v2 — Paper-Guided Optimization
|
| 2 |
+
|
| 3 |
+
Autonomous, research-driven optimization of MoE inference throughput on Apple Silicon.
|
| 4 |
+
|
| 5 |
+
## Objective
|
| 6 |
+
|
| 7 |
+
**Maximize tok/s** on the fixed benchmark while maintaining output quality. Each experiment is guided by a specific technique from a research paper — not random parameter tuning.
|
| 8 |
+
|
| 9 |
+
## Setup
|
| 10 |
+
|
| 11 |
+
1. **Read context files** (in this order):
|
| 12 |
+
- `autoresearch/findings.md` — what previous agents learned. DO NOT repeat failed approaches.
|
| 13 |
+
- `autoresearch/research_queue.md` — prioritized technique queue with paper citations.
|
| 14 |
+
- `autoresearch/seed_papers.md` — curated list of relevant papers.
|
| 15 |
+
- `CLAUDE.md` — architecture overview, constraints, what worked/failed historically.
|
| 16 |
+
- The modular source files (see Files You May Modify below).
|
| 17 |
+
- `autoresearch/benchmark.sh` — the fixed benchmark harness. Do not modify.
|
| 18 |
+
|
| 19 |
+
2. **Create branch**: `git checkout -b autoresearch/<tag>` from `develop`. Propose a tag based on today's date.
|
| 20 |
+
|
| 21 |
+
3. **Verify environment**: Confirm `FLASH_MOE_MODEL` is set and `autoresearch/baseline.txt` exists. If not, tell the human to run `bash autoresearch/prepare.sh`.
|
| 22 |
+
|
| 23 |
+
4. **Run baseline**: `bash autoresearch/benchmark.sh 2>/dev/null` to establish the current starting point.
|
| 24 |
+
|
| 25 |
+
## The Research-Experiment Loop
|
| 26 |
+
|
| 27 |
+
**LOOP FOREVER:**
|
| 28 |
+
|
| 29 |
+
### Step 1: Pick a Technique
|
| 30 |
+
|
| 31 |
+
Read `research_queue.md`. Pick the **top-ranked unchecked technique**. If all are checked, go to Step 1b.
|
| 32 |
+
|
| 33 |
+
### Step 1b: Search for New Papers (every 10 experiments or when queue is empty)
|
| 34 |
+
|
| 35 |
+
Use WebSearch to find recent papers on:
|
| 36 |
+
- "MoE inference optimization 2024 2025"
|
| 37 |
+
- "Apple Silicon Metal GPU kernel optimization"
|
| 38 |
+
- "SSD-based LLM inference offloading"
|
| 39 |
+
- "expert routing prediction mixture of experts"
|
| 40 |
+
- "quantized matrix vector multiply GPU"
|
| 41 |
+
- "speculative decoding sparse models"
|
| 42 |
+
|
| 43 |
+
For each promising paper:
|
| 44 |
+
- Read the abstract and method section (use WebFetch on arxiv/paper URL)
|
| 45 |
+
- Extract the specific technique that could apply to Flash-MoE
|
| 46 |
+
- Add to `research_queue.md` with paper citation and expected impact
|
| 47 |
+
- Rank by: (a) relevance to our bottleneck breakdown, (b) implementation effort, (c) theoretical speedup
|
| 48 |
+
|
| 49 |
+
### Step 2: Read the Paper
|
| 50 |
+
|
| 51 |
+
Use WebSearch/WebFetch to find and read the specific paper section describing the technique. Understand:
|
| 52 |
+
- The exact algorithm or code change
|
| 53 |
+
- Why it works (theoretical basis)
|
| 54 |
+
- What hardware assumptions it makes (may need adaptation for Apple Silicon)
|
| 55 |
+
- Expected speedup and under what conditions
|
| 56 |
+
|
| 57 |
+
### Step 3: Form Hypothesis
|
| 58 |
+
|
| 59 |
+
Write a specific, testable hypothesis:
|
| 60 |
+
```
|
| 61 |
+
HYPOTHESIS: [Technique] from [Paper] should improve tok/s by ~X% because [reason].
|
| 62 |
+
IMPLEMENTATION: Change [specific function/kernel] to [specific change].
|
| 63 |
+
RISK: [What could go wrong — quality, stability, platform-specific issues]
|
| 64 |
+
```
|
| 65 |
+
|
| 66 |
+
### Step 4: Implement
|
| 67 |
+
|
| 68 |
+
Make the change in the appropriate module file(s). Keep changes minimal and focused — one technique per experiment.
|
| 69 |
+
|
| 70 |
+
### Step 5: Build & Benchmark
|
| 71 |
+
|
| 72 |
+
```bash
|
| 73 |
+
git add metal_infer/ && git commit -m "experiment: [technique] (from [paper])"
|
| 74 |
+
bash autoresearch/benchmark.sh 2>/dev/null
|
| 75 |
+
```
|
| 76 |
+
|
| 77 |
+
Parse the BENCH_RESULT line: `tok_s=X.XX math=PASS/FAIL json=PASS/FAIL status=OK/...`
|
| 78 |
+
|
| 79 |
+
### Step 6: Record & Decide
|
| 80 |
+
|
| 81 |
+
Append to `autoresearch/experiments.tsv`:
|
| 82 |
+
```
|
| 83 |
+
id timestamp commit tok_s math json status paper notes
|
| 84 |
+
```
|
| 85 |
+
|
| 86 |
+
**KEEP** if: `tok_s >= previous_best * 0.995` AND both quality gates PASS
|
| 87 |
+
**DISCARD** if: `tok_s < previous_best * 0.995` OR any quality gate FAIL → `git reset --hard HEAD~1`
|
| 88 |
+
|
| 89 |
+
### Step 7: Write Findings
|
| 90 |
+
|
| 91 |
+
Append to `autoresearch/findings.md`:
|
| 92 |
+
```markdown
|
| 93 |
+
## Experiment N: [Technique Name]
|
| 94 |
+
**Paper:** [Title] ([Author], [Year])
|
| 95 |
+
**Hypothesis:** [What we expected and why]
|
| 96 |
+
**Implementation:** [What we changed, 1-2 paragraphs with file/line references]
|
| 97 |
+
**Result:** tok_s=X.XX (baseline: Y.YY, delta: +/-Z%), math=P/F, json=P/F → KEEP/DISCARD
|
| 98 |
+
**Analysis:** [Why it worked or didn't. What we learned about the system.]
|
| 99 |
+
**Next:** [What this result suggests trying next]
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
### Step 8: Update Queue
|
| 103 |
+
|
| 104 |
+
In `research_queue.md`:
|
| 105 |
+
- Check off `[x]` the technique you just tested
|
| 106 |
+
- If the paper suggested related techniques, add them
|
| 107 |
+
- If findings suggest a new direction, add it with rationale
|
| 108 |
+
|
| 109 |
+
### Step 9: Repeat
|
| 110 |
+
|
| 111 |
+
Go to Step 1.
|
| 112 |
+
|
| 113 |
+
## Files You May Modify
|
| 114 |
+
|
| 115 |
+
```
|
| 116 |
+
metal_infer/infer.m — unity build entry point
|
| 117 |
+
metal_infer/config.h — model config, macros
|
| 118 |
+
metal_infer/timing.h — timing globals, telemetry
|
| 119 |
+
metal_infer/fp8.h — FP8 E4M3 support
|
| 120 |
+
metal_infer/weights.h — weight loading, tensor manifest
|
| 121 |
+
metal_infer/cpu_kernels.h — CPU compute kernels, tokenizer
|
| 122 |
+
metal_infer/metal_ctx.h — Metal setup, buffer allocation
|
| 123 |
+
metal_infer/gpu_dispatch.h — GPU dispatch, batched matmul
|
| 124 |
+
metal_infer/expert_io.h — Expert I/O, parallel pread, cache
|
| 125 |
+
metal_infer/layer_forward.h — Forward pass, attention, MoE routing
|
| 126 |
+
metal_infer/generate.h — Generation loop, sampling
|
| 127 |
+
metal_infer/shaders.metal — Metal compute kernels
|
| 128 |
+
```
|
| 129 |
+
|
| 130 |
+
## Files You Must NOT Modify
|
| 131 |
+
|
| 132 |
+
- `autoresearch/benchmark.sh` — the measurement harness is sacred
|
| 133 |
+
- `autoresearch/prepare.sh`
|
| 134 |
+
- `metal_infer/Makefile`
|
| 135 |
+
- `metal_infer/chat.m`, `main.m`, `tokenizer.h`
|
| 136 |
+
- Any Python files (`*.py`)
|
| 137 |
+
- Any iOS app files (`FlashMoE-iOS/`)
|
| 138 |
+
- `CLAUDE.md`
|
| 139 |
+
|
| 140 |
+
## Architecture Constraints — READ CAREFULLY
|
| 141 |
+
|
| 142 |
+
These are hard-won lessons from 58+ prior experiments. Violating them will waste time.
|
| 143 |
+
|
| 144 |
+
### The Unified Memory Constraint
|
| 145 |
+
|
| 146 |
+
On Apple Silicon, **SSD DMA and GPU compute share the same memory controller**. They cannot be profitably overlapped. The GPU's dequant kernels are bandwidth-saturated at ~418 GiB/s. Even small background SSD DMA causes disproportionate GPU latency spikes.
|
| 147 |
+
|
| 148 |
+
**Do NOT attempt:**
|
| 149 |
+
- Overlapping SSD reads with GPU compute (validated: net negative)
|
| 150 |
+
- dispatch_io (70% slower)
|
| 151 |
+
- F_RDADVISE / speculative prefetch (net 0%)
|
| 152 |
+
- mmap for expert files (5x slower)
|
| 153 |
+
|
| 154 |
+
### The Caching Constraint
|
| 155 |
+
|
| 156 |
+
**Trust the OS page cache.** Every custom expert cache we tried was slower:
|
| 157 |
+
- Metal LRU cache: -38%
|
| 158 |
+
- malloc cache: -20%
|
| 159 |
+
- LZ4 compressed cache: -13%
|
| 160 |
+
- Speculative early routing: -38%
|
| 161 |
+
|
| 162 |
+
### GPU Constraints
|
| 163 |
+
|
| 164 |
+
- Metal GPU is **memory-bandwidth-bound** at ~418 GiB/s
|
| 165 |
+
- No spin-polling (CPU thermal competes with GPU)
|
| 166 |
+
- Expert files are ~6.75MB each; NVMe doesn't care about scatter at this granularity
|
| 167 |
+
|
| 168 |
+
### What Already Works Well (don't reinvent)
|
| 169 |
+
|
| 170 |
+
- FMA dequant kernel: `fma(nibble, scale*x, bias*x)`
|
| 171 |
+
- Deferred CMD3: expert forward pass submitted without waiting
|
| 172 |
+
- BLAS delta-net: Accelerate framework for 64-head state recurrence
|
| 173 |
+
- GCD parallel pread: dispatch groups for K experts
|
| 174 |
+
- GPU fused attention: RoPE + QK norm fused
|
| 175 |
+
- Fused moe_combine_residual: single kernel for combine + residual + sigmoid gate
|
| 176 |
+
- CMD1+CMD2 merge: fewer command buffers for linear attention layers
|
| 177 |
+
- Delta-net kernel fusion: merged pass 2+3 in gated_delta_net_step
|
| 178 |
+
- FP16 accumulation kernels: half-precision variants of all dequant kernels
|
| 179 |
+
- FP8 KV cache: E4M3 quantized attention cache
|
| 180 |
+
|
| 181 |
+
### Bottleneck Breakdown (397B, 4-bit, M3 Max)
|
| 182 |
+
|
| 183 |
+
| Phase | Time/Layer | % Total | Notes |
|
| 184 |
+
|-------|-----------|---------|-------|
|
| 185 |
+
| Expert I/O (pread) | 2.41ms | 56% | THE bottleneck |
|
| 186 |
+
| GPU compute (CMD1+CMD2+CMD3) | 1.81ms | 42% | Bandwidth-saturated |
|
| 187 |
+
| CPU routing | 0.003ms | <1% | Negligible |
|
| 188 |
+
| Misc overhead | 0.05ms | 1% | Command encoding |
|
| 189 |
+
|
| 190 |
+
## Decision Rules
|
| 191 |
+
|
| 192 |
+
- **Keep** if: `tok_s >= previous_best * 0.995` AND both quality gates PASS
|
| 193 |
+
- **Discard** if: `tok_s < previous_best * 0.995` OR any quality gate FAIL
|
| 194 |
+
- If **3 consecutive discards** with similar approaches: move to a different area
|
| 195 |
+
- If **build fails**: quick fix (2 tries max), then revert
|
| 196 |
+
- After every **keep**: update baseline to new tok_s
|
| 197 |
+
|
| 198 |
+
## Simplicity Criterion
|
| 199 |
+
|
| 200 |
+
Simpler is better. A 0.5% improvement that adds 50 lines? Borderline. A 0.5% improvement from *removing* code? Definitely keep. The goal is a lean, fast engine.
|
| 201 |
+
|
| 202 |
+
## Safety
|
| 203 |
+
|
| 204 |
+
- Primary dev machine with 48GB unified RAM
|
| 205 |
+
- Do not allocate more than ~200MB of new Metal buffers
|
| 206 |
+
- Do not create files larger than 10MB
|
| 207 |
+
- Do not modify the build system or add dependencies
|
| 208 |
+
- If build fails, revert immediately
|
| 209 |
+
|
| 210 |
+
## NEVER STOP
|
| 211 |
+
|
| 212 |
+
Once the loop begins, do NOT pause to ask the human. The human may be asleep. Run autonomously until manually interrupted. If you run out of queued techniques, search for more papers. Each experiment takes ~2-5 minutes, so you can run ~12-30 per hour. The loop runs until the human stops you.
|
autoresearch/research_queue.md
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Research Queue
|
| 2 |
+
|
| 3 |
+
Prioritized techniques to test, extracted from papers. Checked items have been tested — see `findings.md` for results.
|
| 4 |
+
|
| 5 |
+
## Priority 1 — Expert I/O Optimization (56% of bottleneck)
|
| 6 |
+
|
| 7 |
+
- [ ] **Frequency-based expert pinning** — Keep the top-N most frequently activated experts in a persistent mmap'd buffer. OS page cache does LRU, but frequency-based pinning is different (sticky hot set). From "LLM in a Flash" (Alizadeh et al., Apple, 2023). Risk: we tested custom caches before and they were slower, but pinning != caching.
|
| 8 |
+
- [ ] **Expert co-activation clustering** — Reorder expert indices within each layer file so frequently co-activated experts are physically adjacent. Sequential NVMe reads are faster than scattered reads at small granularity. From "MoE-Infinity" (Xue et al., 2024). Note: we tested file-level clustering at 7MB granularity (0% improvement), but sub-expert byte-level adjacency within a single file is untested.
|
| 9 |
+
- [ ] **Attention-to-routing correlation prefetch** — Use layer N's attention output distribution to predict layer N+1's expert routing before the actual router MLP runs. Gives a 1-2ms head start on pread. From "Pre-gated MoE" (Hwang et al., 2024). Risk: prediction accuracy may be too low (our temporal predictor hit 25%).
|
| 10 |
+
- [ ] **NVMe page-aligned expert reads** — Ensure pread offsets and sizes are 4KB-aligned to avoid partial page reads. From "FlashNeuron" (Song et al., 2021). Quick check: are our current expert offsets already aligned?
|
| 11 |
+
- [ ] **Compressed expert transfer** — Use hardware-accelerated LZ4/zstd decompression (Accelerate.framework) on expert weights during SSD→memory transfer. Different from our prior test (software LZ4 was -13%) because hardware decompression has near-zero CPU cost. Needs investigation of Accelerate's compression APIs.
|
| 12 |
+
|
| 13 |
+
## Priority 2 — GPU Compute Optimization (30% of bottleneck)
|
| 14 |
+
|
| 15 |
+
- [ ] **Register-tiled GEMV** — Process multiple output rows per thread (thread coarsening) to amortize shared memory loads. Each thread computes 2-4 output elements instead of 1. From "AWQ" (Lin et al., 2023) and general GPU optimization literature. Currently ROWS_PER_TG=8 with 1 row per SIMD group — could do 2 rows per SIMD group.
|
| 16 |
+
- [ ] **Persistent threadgroups** — Keep threadgroups alive across multiple matvec dispatches within a command buffer using `[[threadgroup_size_in_memory]]`. Avoids threadgroup allocation/deallocation overhead. From CUTLASS (NVIDIA, 2022) — Metal equivalent is `dispatchThreadgroups:threadsPerThreadgroup:` with persistent launch.
|
| 17 |
+
- [ ] **SIMD shuffle reduction** — Replace `threadgroup float shared[32]` reductions with `simd_shuffle_xor` cascade. Avoids shared memory roundtrip for the final reduction step. From "FlashDecoding++" (Hong et al., 2024).
|
| 18 |
+
- [ ] **Mixed-precision with Kahan compensated summation** — Accumulate in fp16 for speed but use Kahan summation to recover fp32 accuracy. Best of both worlds: fp16 throughput with fp32 error bounds. From numerical analysis literature.
|
| 19 |
+
- [ ] **Vectorized nibble extraction** — Use `as_type<uchar4>()` to extract 4 nibbles simultaneously instead of sequential shift+mask. May improve instruction-level parallelism. From GPU programming guides.
|
| 20 |
+
|
| 21 |
+
## Priority 3 — Pipeline Optimization (10% of bottleneck)
|
| 22 |
+
|
| 23 |
+
- [ ] **Background command buffer encoding** — Encode CMD1 for layer N+1 on a background thread while GPU executes CMD3 for layer N. GCD dispatch_async + MTLCommandBuffer from shared queue. From general Metal best practices.
|
| 24 |
+
- [ ] **Layer-wise adaptive K** — Not all layers need the same number of experts. Early layers may work with K=2 while later layers need K=4. Measure per-layer routing entropy to auto-select K. From "Adaptive Computation in MoE" (various, 2024).
|
| 25 |
+
- [ ] **Speculative token execution** — Generate 2 candidate next tokens, run both through the first few layers, pick the winner early. For MoE, expert I/O scales per-token so this is risky. From "Sequoia" (Chen et al., 2024). Note: we tested MTP speculative decoding before (break-even) — this is a different approach.
|
| 26 |
+
|
| 27 |
+
## Priority 4 — Novel Techniques (from paper search)
|
| 28 |
+
|
| 29 |
+
- [ ] **Layer circuit duplication (adaptive depth)** — Repeat specific contiguous layer blocks (e.g. layers 20-25) during inference for harder prompts without any extra weights. "Functional circuits" across consecutive layers amplify reasoning ability. Use math/logic probes as heatmaps to identify which layer blocks to repeat. For simple prompts run once, for complex prompts run the circuit twice. Zero memory overhead. From DNH blog post (2025) on dense transformer layer duplication — needs adaptation for hybrid GatedDeltaNet + full attention architecture. Our 30 linear + 10 full attention layers may have different circuit patterns. Key question: does delta-net state accumulation remain stable through repeated layers?
|
| 30 |
+
- [ ] **Workload-aware page cache priming** — Use calibration-derived expert activation mass data to `madvise(MADV_WILLNEED)` on the hottest expert files at startup, improving cold-start cache hit rate from ~71% to potentially 85%+. From REAP-swap (Sero et al., 2025). Low effort (~20 lines), needs freq_data.json from calibration run.
|
| 31 |
+
- [ ] **MSE-optimal clipping for 2-bit** — Search 20 clipping ratios (0.7-1.0 of weight range) to minimize reconstruction MSE instead of using raw min/max. 15-30% RMSE reduction at 2-bit. May fix JSON corruption. From JANG (2025). Already implemented in `repack_experts_2bit.py --mse-clip` — needs benchmarking.
|
| 32 |
+
- [ ] **3-bit expert quantization** — Sweet spot between 2-bit (broken JSON) and 4-bit (209GB). 25% less I/O per expert with likely good-enough quality. Requires new Metal kernel for cross-byte-boundary bit extraction. From JANG adaptive quantization profiles.
|
| 33 |
+
- [ ] **DWQ/GPTQ error compensation** — Column-by-column Hessian-guided quantization that distributes rounding errors optimally. Makes 2-bit usable for JSON/tool calling. Already implemented in `gptq_requantize.py` — needs end-to-end validation with tool calling test.
|
| 34 |
+
- [ ] **Kahan compensated fp16 accumulation** — Use Kahan summation with fp16 accumulators to get fp16 throughput with fp32 accuracy bounds. Different from our current fp16 toggle which has no error compensation. From numerical analysis literature.
|
autoresearch/seed_papers.md
ADDED
|
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Seed Papers for Flash-MoE Optimization
|
| 2 |
+
|
| 3 |
+
Pre-curated papers relevant to our specific architecture (MoE + SSD streaming + Metal GPU + Apple Silicon). The autoresearch agent should search for and read these, plus discover new papers via WebSearch.
|
| 4 |
+
|
| 5 |
+
## Expert I/O & Offloading
|
| 6 |
+
|
| 7 |
+
1. **"LLM in a Flash: Efficient Large Language Model Inference with Limited Memory"**
|
| 8 |
+
- Authors: Alizadeh et al. (Apple), 2023
|
| 9 |
+
- URL: https://arxiv.org/abs/2312.11514
|
| 10 |
+
- Relevance: THE foundational paper for our approach. Windowing strategy, row-column bundling for flash reads, preloading based on activation sparsity. We already use their core idea but may have missed specific techniques.
|
| 11 |
+
|
| 12 |
+
2. **"MoE-Infinity: Offloading-Efficient MoE Model Serving"**
|
| 13 |
+
- Authors: Xue et al., 2024
|
| 14 |
+
- URL: https://arxiv.org/abs/2401.14361
|
| 15 |
+
- Relevance: Expert-level offloading with activation-aware prefetch. Their "expert popularity" tracking could improve our page cache hit rate.
|
| 16 |
+
|
| 17 |
+
3. **"Pre-gated MoE: An Algorithm-System Co-Design for Fast and Scalable Mixture-of-Expert Inference"**
|
| 18 |
+
- Authors: Hwang et al., 2024
|
| 19 |
+
- URL: https://arxiv.org/abs/2308.12066
|
| 20 |
+
- Relevance: Predicts expert routing from the previous layer's output, enabling prefetch. Different from our temporal predictor (which failed at 25% accuracy).
|
| 21 |
+
|
| 22 |
+
4. **"PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU"**
|
| 23 |
+
- Authors: Song et al., 2023
|
| 24 |
+
- URL: https://arxiv.org/abs/2312.12456
|
| 25 |
+
- Relevance: Hot/cold neuron partitioning between GPU and CPU. Our tiered quantization is related but PowerInfer's activation prediction is more sophisticated.
|
| 26 |
+
|
| 27 |
+
## GPU Kernel Optimization
|
| 28 |
+
|
| 29 |
+
5. **"AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration"**
|
| 30 |
+
- Authors: Lin et al., 2023
|
| 31 |
+
- URL: https://arxiv.org/abs/2306.00978
|
| 32 |
+
- Relevance: Activation-aware scaling before quantization + optimized GEMV kernels. Their kernel tiling strategy may apply to our dequant matvec.
|
| 33 |
+
|
| 34 |
+
6. **"FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning"**
|
| 35 |
+
- Authors: Dao, 2023
|
| 36 |
+
- URL: https://arxiv.org/abs/2307.08691
|
| 37 |
+
- Relevance: Online softmax, warp specialization, better tiling. We implemented fused online softmax based on this but there may be Metal-specific optimizations we missed.
|
| 38 |
+
|
| 39 |
+
7. **"FlashDecoding++: Faster Large Language Model Inference with Asynchronous Softmax"**
|
| 40 |
+
- Authors: Hong et al., 2024
|
| 41 |
+
- URL: https://arxiv.org/abs/2311.01282
|
| 42 |
+
- Relevance: Flat GEMV optimization for decode phase. Their "unified maximum" technique eliminates the softmax sync point.
|
| 43 |
+
|
| 44 |
+
8. **"QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving"**
|
| 45 |
+
- Authors: Lin et al., 2024
|
| 46 |
+
- URL: https://arxiv.org/abs/2405.04532
|
| 47 |
+
- Relevance: W4A8 serving with progressive quantization. Their SmoothAttention technique for KV4 is directly applicable to our FP8 KV cache.
|
| 48 |
+
|
| 49 |
+
## Speculative & Predictive Techniques
|
| 50 |
+
|
| 51 |
+
9. **"Sequoia: Scalable and Robust Speculative Decoding"**
|
| 52 |
+
- Authors: Chen et al., 2024
|
| 53 |
+
- URL: https://arxiv.org/abs/2402.12374
|
| 54 |
+
- Relevance: Speculative decoding for offloaded models. Our MTP test was break-even, but Sequoia's tree-based approach handles MoE differently.
|
| 55 |
+
|
| 56 |
+
10. **"Mixtral of Experts"**
|
| 57 |
+
- Authors: Jiang et al. (Mistral AI), 2024
|
| 58 |
+
- URL: https://arxiv.org/abs/2401.04088
|
| 59 |
+
- Relevance: The MoE architecture paper. Details on expert routing behavior, load balancing, and routing statistics that inform prefetch strategies.
|
| 60 |
+
|
| 61 |
+
## Apple Silicon Specific
|
| 62 |
+
|
| 63 |
+
11. **"Accelerating Large Language Model Decoding with Speculative Sampling"**
|
| 64 |
+
- Authors: Leviathan et al., 2022
|
| 65 |
+
- URL: https://arxiv.org/abs/2302.01318
|
| 66 |
+
- Relevance: Original speculative decoding paper. The acceptance/rejection scheme could work with our tiered approach.
|
| 67 |
+
|
| 68 |
+
12. **Apple Metal Best Practices Guide**
|
| 69 |
+
- URL: https://developer.apple.com/library/archive/documentation/3DDrawing/Conceptual/MTLBestPracticesGuide/
|
| 70 |
+
- Relevance: Official guidance on command buffer management, resource allocation patterns, and GPU occupancy on Apple Silicon.
|
| 71 |
+
|
| 72 |
+
## To Discover
|
| 73 |
+
|
| 74 |
+
The agent should search for papers published in 2025-2026 on:
|
| 75 |
+
- "mixture of experts inference optimization"
|
| 76 |
+
- "apple silicon gpu compute optimization"
|
| 77 |
+
- "SSD offloading large language models"
|
| 78 |
+
- "4-bit quantized matmul GPU kernel"
|
| 79 |
+
- "expert routing prediction"
|
build_expert_index.py
ADDED
|
@@ -0,0 +1,167 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Build expert_index.json from a model's safetensors index.
|
| 3 |
+
|
| 4 |
+
Scans model.safetensors.index.json to find all expert (switch_mlp) weight tensors
|
| 5 |
+
and computes their byte offsets and strides for the repacking script.
|
| 6 |
+
|
| 7 |
+
Usage:
|
| 8 |
+
python build_expert_index.py --model ~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import argparse
|
| 12 |
+
import json
|
| 13 |
+
import os
|
| 14 |
+
import struct
|
| 15 |
+
import sys
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
from collections import defaultdict
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def parse_safetensors_header(filepath):
|
| 21 |
+
"""Parse a safetensors file header. Returns (header_dict, data_start_offset)."""
|
| 22 |
+
with open(filepath, 'rb') as f:
|
| 23 |
+
header_len = struct.unpack('<Q', f.read(8))[0]
|
| 24 |
+
header = json.loads(f.read(header_len))
|
| 25 |
+
data_start = 8 + header_len
|
| 26 |
+
return header, data_start
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def main():
|
| 30 |
+
parser = argparse.ArgumentParser(description='Build expert_index.json from safetensors')
|
| 31 |
+
parser.add_argument('--model', type=str,
|
| 32 |
+
default=os.path.expanduser(
|
| 33 |
+
'~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit'),
|
| 34 |
+
help='Path to model directory')
|
| 35 |
+
parser.add_argument('--output', type=str, default='expert_index.json',
|
| 36 |
+
help='Output path for expert_index.json')
|
| 37 |
+
args = parser.parse_args()
|
| 38 |
+
|
| 39 |
+
model_path = Path(args.model)
|
| 40 |
+
|
| 41 |
+
# Find the snapshot directory if it exists (HF cache layout)
|
| 42 |
+
# Check for model.safetensors.index.json in the model path or its snapshots
|
| 43 |
+
index_path = model_path / 'model.safetensors.index.json'
|
| 44 |
+
if not index_path.exists():
|
| 45 |
+
# Try snapshot subdirectories
|
| 46 |
+
snapshots_dir = model_path / 'snapshots'
|
| 47 |
+
if snapshots_dir.exists():
|
| 48 |
+
for snap in sorted(snapshots_dir.iterdir()):
|
| 49 |
+
candidate = snap / 'model.safetensors.index.json'
|
| 50 |
+
if candidate.exists():
|
| 51 |
+
model_path = snap
|
| 52 |
+
index_path = candidate
|
| 53 |
+
break
|
| 54 |
+
|
| 55 |
+
if not index_path.exists():
|
| 56 |
+
print(f"ERROR: {index_path} not found", file=sys.stderr)
|
| 57 |
+
sys.exit(1)
|
| 58 |
+
|
| 59 |
+
print(f"Model path: {model_path}")
|
| 60 |
+
print(f"Index: {index_path}")
|
| 61 |
+
|
| 62 |
+
with open(index_path) as f:
|
| 63 |
+
idx = json.load(f)
|
| 64 |
+
|
| 65 |
+
weight_map = idx['weight_map']
|
| 66 |
+
|
| 67 |
+
# Find all expert tensor names: pattern is
|
| 68 |
+
# model.layers.{L}.switch_mlp.{gate_proj|up_proj|down_proj}.{weight|scales|biases}
|
| 69 |
+
# or with language_model. prefix
|
| 70 |
+
import re
|
| 71 |
+
expert_pattern = re.compile(
|
| 72 |
+
r'(?:language_model\.)?model\.layers\.(\d+)\.(?:mlp\.)?switch_mlp\.'
|
| 73 |
+
r'(gate_proj|up_proj|down_proj)\.(weight|scales|biases)$'
|
| 74 |
+
)
|
| 75 |
+
|
| 76 |
+
# Group expert tensors by layer and component
|
| 77 |
+
layer_tensors = defaultdict(dict) # layer_idx -> {component_key -> tensor_name}
|
| 78 |
+
for name in weight_map:
|
| 79 |
+
m = expert_pattern.match(name)
|
| 80 |
+
if m:
|
| 81 |
+
layer_idx = int(m.group(1))
|
| 82 |
+
proj = m.group(2) # gate_proj, up_proj, down_proj
|
| 83 |
+
part = m.group(3) # weight, scales, biases
|
| 84 |
+
component_key = f"{proj}.{part}"
|
| 85 |
+
layer_tensors[layer_idx][component_key] = name
|
| 86 |
+
|
| 87 |
+
if not layer_tensors:
|
| 88 |
+
print("ERROR: No expert tensors found in weight_map", file=sys.stderr)
|
| 89 |
+
print("Sample tensor names:", list(weight_map.keys())[:20], file=sys.stderr)
|
| 90 |
+
sys.exit(1)
|
| 91 |
+
|
| 92 |
+
num_layers = len(layer_tensors)
|
| 93 |
+
print(f"Found expert tensors in {num_layers} layers")
|
| 94 |
+
|
| 95 |
+
# Parse safetensors headers to get exact offsets
|
| 96 |
+
header_cache = {}
|
| 97 |
+
expert_reads = {}
|
| 98 |
+
|
| 99 |
+
for layer_idx in sorted(layer_tensors.keys()):
|
| 100 |
+
components = layer_tensors[layer_idx]
|
| 101 |
+
layer_reads = {}
|
| 102 |
+
|
| 103 |
+
for comp_key, tensor_name in sorted(components.items()):
|
| 104 |
+
shard_file = weight_map[tensor_name]
|
| 105 |
+
filepath = model_path / shard_file
|
| 106 |
+
|
| 107 |
+
if shard_file not in header_cache:
|
| 108 |
+
header_cache[shard_file] = parse_safetensors_header(str(filepath))
|
| 109 |
+
|
| 110 |
+
header, data_start = header_cache[shard_file]
|
| 111 |
+
|
| 112 |
+
if tensor_name not in header:
|
| 113 |
+
# Try without language_model prefix
|
| 114 |
+
alt_name = tensor_name
|
| 115 |
+
if alt_name.startswith("language_model."):
|
| 116 |
+
alt_name = alt_name[len("language_model."):]
|
| 117 |
+
if alt_name not in header:
|
| 118 |
+
print(f"WARNING: {tensor_name} not in {shard_file} header, skipping")
|
| 119 |
+
continue
|
| 120 |
+
tensor_name_in_header = alt_name
|
| 121 |
+
else:
|
| 122 |
+
tensor_name_in_header = tensor_name
|
| 123 |
+
|
| 124 |
+
meta = header[tensor_name_in_header]
|
| 125 |
+
offsets = meta['data_offsets']
|
| 126 |
+
shape = meta['shape']
|
| 127 |
+
total_size = offsets[1] - offsets[0]
|
| 128 |
+
abs_offset = data_start + offsets[0]
|
| 129 |
+
|
| 130 |
+
# For expert tensors, the first dimension is num_experts (256)
|
| 131 |
+
# Expert stride = total_size / num_experts
|
| 132 |
+
num_experts = shape[0]
|
| 133 |
+
expert_size = total_size // num_experts
|
| 134 |
+
|
| 135 |
+
layer_reads[comp_key] = {
|
| 136 |
+
"file": shard_file,
|
| 137 |
+
"abs_offset": abs_offset,
|
| 138 |
+
"expert_stride": expert_size,
|
| 139 |
+
"expert_size": expert_size,
|
| 140 |
+
"total_size": total_size,
|
| 141 |
+
"shape": shape,
|
| 142 |
+
"dtype": meta['dtype'],
|
| 143 |
+
}
|
| 144 |
+
|
| 145 |
+
expert_reads[str(layer_idx)] = layer_reads
|
| 146 |
+
if layer_idx == 0:
|
| 147 |
+
print(f"\nLayer 0 components:")
|
| 148 |
+
for k, v in sorted(layer_reads.items()):
|
| 149 |
+
print(f" {k}: shape={v['shape']}, expert_size={v['expert_size']}, "
|
| 150 |
+
f"file={v['file']}")
|
| 151 |
+
|
| 152 |
+
# Build output
|
| 153 |
+
output = {
|
| 154 |
+
"model_path": str(model_path),
|
| 155 |
+
"expert_reads": expert_reads,
|
| 156 |
+
}
|
| 157 |
+
|
| 158 |
+
with open(args.output, 'w') as f:
|
| 159 |
+
json.dump(output, f, indent=2)
|
| 160 |
+
|
| 161 |
+
print(f"\nWrote {args.output}")
|
| 162 |
+
print(f"Layers: {num_layers}")
|
| 163 |
+
print(f"Components per layer: {len(next(iter(expert_reads.values())))}")
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
if __name__ == '__main__':
|
| 167 |
+
main()
|
cloud_repack.py
ADDED
|
@@ -0,0 +1,1355 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Cloud repack pipeline for Flash-MoE models.
|
| 3 |
+
|
| 4 |
+
Downloads an MLX model from HuggingFace, repacks expert weights into
|
| 5 |
+
per-layer binary files, extracts non-expert weights (with optional split
|
| 6 |
+
for iOS Metal 4GB limit), and uploads the result to a new HF repo.
|
| 7 |
+
|
| 8 |
+
Supports three modes:
|
| 9 |
+
- 4bit (default): All experts at 4-bit quantization.
|
| 10 |
+
- tiered: Hot experts at 4-bit, cold experts requantized to 2-bit (MSE-optimal).
|
| 11 |
+
- gptq: GPTQ error-compensated 2-bit (placeholder, requires calibration data).
|
| 12 |
+
|
| 13 |
+
Runs on Modal (https://modal.com) — no local disk or GPU needed.
|
| 14 |
+
|
| 15 |
+
Usage:
|
| 16 |
+
# Install modal client
|
| 17 |
+
pip install modal
|
| 18 |
+
modal setup # one-time auth
|
| 19 |
+
|
| 20 |
+
# 4-bit repack (default)
|
| 21 |
+
modal run cloud_repack.py --source mlx-community/Qwen3.5-122B-A10B-4bit \
|
| 22 |
+
--dest alexintosh/Qwen3.5-122B-A10B-Q4-FlashMoE \
|
| 23 |
+
--split 3.5
|
| 24 |
+
|
| 25 |
+
# Tiered repack (top 20% hot at 4-bit, rest at 2-bit)
|
| 26 |
+
modal run cloud_repack.py --source mlx-community/Qwen3.5-35B-A3B-4bit \
|
| 27 |
+
--dest alexintosh/Qwen3.5-35B-A3B-Q4-Tiered-FlashMoE \
|
| 28 |
+
--mode tiered --hot-ratio 0.2 --split 3.5
|
| 29 |
+
|
| 30 |
+
# Repack 397B (needs large volume)
|
| 31 |
+
modal run cloud_repack.py --source mlx-community/Qwen3.5-397B-A17B-4bit \
|
| 32 |
+
--dest alexintosh/Qwen3.5-397B-A17B-Q4-FlashMoE \
|
| 33 |
+
--split 3.5 --volume-size 500
|
| 34 |
+
"""
|
| 35 |
+
|
| 36 |
+
import modal
|
| 37 |
+
import os
|
| 38 |
+
from pathlib import Path
|
| 39 |
+
|
| 40 |
+
app = modal.App("flash-moe-repack")
|
| 41 |
+
|
| 42 |
+
# Volume for storing model data during repacking
|
| 43 |
+
vol = modal.Volume.from_name("flash-moe-repack-vol", create_if_missing=True)
|
| 44 |
+
|
| 45 |
+
image = (
|
| 46 |
+
modal.Image.debian_slim(python_version="3.11")
|
| 47 |
+
.pip_install(
|
| 48 |
+
"huggingface_hub",
|
| 49 |
+
"safetensors",
|
| 50 |
+
"numpy",
|
| 51 |
+
"torch", # needed for safetensors loading
|
| 52 |
+
"scipy", # for future GPTQ support
|
| 53 |
+
)
|
| 54 |
+
)
|
| 55 |
+
|
| 56 |
+
# Higher-resource image config for tiered repacking (needs more RAM for requantization)
|
| 57 |
+
tiered_image = image
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
@app.function(
|
| 61 |
+
image=image,
|
| 62 |
+
volumes={"/data": vol},
|
| 63 |
+
secrets=[modal.Secret.from_name("huggingface-secret", required_keys=["HF_TOKEN"])],
|
| 64 |
+
timeout=7200, # 2 hours max
|
| 65 |
+
memory=16384, # 16 GB RAM
|
| 66 |
+
cpu=4,
|
| 67 |
+
)
|
| 68 |
+
def repack_model(source: str, dest: str, split_gb: float = 0, hf_token: str = ""):
|
| 69 |
+
"""Download, repack, and upload a Flash-MoE model."""
|
| 70 |
+
import json
|
| 71 |
+
import struct
|
| 72 |
+
import time
|
| 73 |
+
from pathlib import Path
|
| 74 |
+
from collections import defaultdict
|
| 75 |
+
from huggingface_hub import HfApi, snapshot_download
|
| 76 |
+
|
| 77 |
+
WORK = Path("/data/work")
|
| 78 |
+
MODEL_DIR = WORK / "model"
|
| 79 |
+
OUTPUT_DIR = WORK / "output"
|
| 80 |
+
|
| 81 |
+
# Clean previous runs
|
| 82 |
+
import shutil
|
| 83 |
+
if OUTPUT_DIR.exists():
|
| 84 |
+
shutil.rmtree(OUTPUT_DIR)
|
| 85 |
+
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
| 86 |
+
|
| 87 |
+
# ----------------------------------------------------------------
|
| 88 |
+
# Step 1: Download model from HuggingFace
|
| 89 |
+
# ----------------------------------------------------------------
|
| 90 |
+
print(f"\n{'='*60}")
|
| 91 |
+
print(f"Step 1: Downloading {source}")
|
| 92 |
+
print(f"{'='*60}")
|
| 93 |
+
|
| 94 |
+
t0 = time.time()
|
| 95 |
+
model_path = Path(snapshot_download(
|
| 96 |
+
source,
|
| 97 |
+
local_dir=str(MODEL_DIR),
|
| 98 |
+
token=hf_token or None,
|
| 99 |
+
))
|
| 100 |
+
print(f"Downloaded in {time.time()-t0:.0f}s to {model_path}")
|
| 101 |
+
|
| 102 |
+
# ----------------------------------------------------------------
|
| 103 |
+
# Step 2: Parse config
|
| 104 |
+
# ----------------------------------------------------------------
|
| 105 |
+
print(f"\n{'='*60}")
|
| 106 |
+
print(f"Step 2: Parsing model config")
|
| 107 |
+
print(f"{'='*60}")
|
| 108 |
+
|
| 109 |
+
config = json.load(open(model_path / "config.json"))
|
| 110 |
+
tc = config.get("text_config", config)
|
| 111 |
+
|
| 112 |
+
num_layers = tc["num_hidden_layers"]
|
| 113 |
+
num_experts = tc["num_experts"]
|
| 114 |
+
hidden_size = tc["hidden_size"]
|
| 115 |
+
moe_intermediate = tc["moe_intermediate_size"]
|
| 116 |
+
group_size = config.get("quantization", {}).get("group_size", 64)
|
| 117 |
+
bits = config.get("quantization", {}).get("bits", 4)
|
| 118 |
+
|
| 119 |
+
print(f" Layers: {num_layers}")
|
| 120 |
+
print(f" Experts: {num_experts}")
|
| 121 |
+
print(f" Hidden: {hidden_size}")
|
| 122 |
+
print(f" MoE intermediate: {moe_intermediate}")
|
| 123 |
+
print(f" Quantization: {bits}-bit, group_size={group_size}")
|
| 124 |
+
|
| 125 |
+
# Compute expert size
|
| 126 |
+
vals_per_u32 = 32 // bits
|
| 127 |
+
# gate_proj / up_proj: [moe_intermediate, hidden_size]
|
| 128 |
+
gate_w = moe_intermediate * ((hidden_size + vals_per_u32 - 1) // vals_per_u32) * 4
|
| 129 |
+
gate_s = moe_intermediate * ((hidden_size + group_size - 1) // group_size) * 2
|
| 130 |
+
gate_b = gate_s
|
| 131 |
+
# down_proj: [hidden_size, moe_intermediate]
|
| 132 |
+
down_w = hidden_size * ((moe_intermediate + vals_per_u32 - 1) // vals_per_u32) * 4
|
| 133 |
+
down_s = hidden_size * ((moe_intermediate + group_size - 1) // group_size) * 2
|
| 134 |
+
down_b = down_s
|
| 135 |
+
|
| 136 |
+
expert_size = 3 * (gate_w + gate_s + gate_b) # gate + up (same shape) + down (same total)
|
| 137 |
+
# Actually: gate + up have same layout, down is different
|
| 138 |
+
expert_size = 2 * (gate_w + gate_s + gate_b) + (down_w + down_s + down_b)
|
| 139 |
+
|
| 140 |
+
print(f" Expert size: {expert_size} bytes ({expert_size/1e6:.2f} MB)")
|
| 141 |
+
print(f" Total expert data: {expert_size * num_experts * num_layers / 1e9:.1f} GB")
|
| 142 |
+
|
| 143 |
+
# Component layout within each expert block
|
| 144 |
+
components = []
|
| 145 |
+
off = 0
|
| 146 |
+
for proj, rows, cols in [("gate_proj", moe_intermediate, hidden_size),
|
| 147 |
+
("up_proj", moe_intermediate, hidden_size),
|
| 148 |
+
("down_proj", hidden_size, moe_intermediate)]:
|
| 149 |
+
w_size = rows * ((cols + vals_per_u32 - 1) // vals_per_u32) * 4
|
| 150 |
+
s_size = rows * ((cols + group_size - 1) // group_size) * 2
|
| 151 |
+
b_size = s_size
|
| 152 |
+
components.append((f"{proj}.weight", off, w_size)); off += w_size
|
| 153 |
+
components.append((f"{proj}.scales", off, s_size)); off += s_size
|
| 154 |
+
components.append((f"{proj}.biases", off, b_size)); off += b_size
|
| 155 |
+
|
| 156 |
+
assert off == expert_size, f"Component layout mismatch: {off} != {expert_size}"
|
| 157 |
+
|
| 158 |
+
# ----------------------------------------------------------------
|
| 159 |
+
# Step 3: Build expert index from safetensors
|
| 160 |
+
# ----------------------------------------------------------------
|
| 161 |
+
print(f"\n{'='*60}")
|
| 162 |
+
print(f"Step 3: Building expert index")
|
| 163 |
+
print(f"{'='*60}")
|
| 164 |
+
|
| 165 |
+
index_file = model_path / "model.safetensors.index.json"
|
| 166 |
+
with open(index_file) as f:
|
| 167 |
+
idx = json.load(f)
|
| 168 |
+
weight_map = idx["weight_map"]
|
| 169 |
+
|
| 170 |
+
# Parse safetensors headers
|
| 171 |
+
header_cache = {}
|
| 172 |
+
def get_header(filename):
|
| 173 |
+
if filename not in header_cache:
|
| 174 |
+
filepath = model_path / filename
|
| 175 |
+
with open(filepath, 'rb') as f:
|
| 176 |
+
header_len = struct.unpack('<Q', f.read(8))[0]
|
| 177 |
+
header = json.loads(f.read(header_len))
|
| 178 |
+
data_start = 8 + header_len
|
| 179 |
+
header_cache[filename] = (header, data_start)
|
| 180 |
+
return header_cache[filename]
|
| 181 |
+
|
| 182 |
+
# Build expert reads: for each layer/expert/component, record file + offset + size
|
| 183 |
+
import re
|
| 184 |
+
expert_pattern = re.compile(
|
| 185 |
+
r'language_model\.model\.layers\.(\d+)\.mlp\.switch_mlp\.'
|
| 186 |
+
r'(gate_proj|up_proj|down_proj)\.(weight|scales|biases)'
|
| 187 |
+
)
|
| 188 |
+
|
| 189 |
+
expert_reads = {} # (layer, expert_idx, component_name) -> (file, offset, size)
|
| 190 |
+
|
| 191 |
+
for tensor_name, filename in weight_map.items():
|
| 192 |
+
m = expert_pattern.match(tensor_name)
|
| 193 |
+
if not m:
|
| 194 |
+
continue
|
| 195 |
+
layer = int(m.group(1))
|
| 196 |
+
proj = m.group(2)
|
| 197 |
+
attr = m.group(3)
|
| 198 |
+
comp_name = f"{proj}.{attr}"
|
| 199 |
+
|
| 200 |
+
header, data_start = get_header(filename)
|
| 201 |
+
meta = header[tensor_name]
|
| 202 |
+
t_off = meta['data_offsets']
|
| 203 |
+
byte_len = t_off[1] - t_off[0]
|
| 204 |
+
shape = meta['shape']
|
| 205 |
+
|
| 206 |
+
# shape[0] = num_experts, rest = per-expert dims
|
| 207 |
+
per_expert_size = byte_len // shape[0]
|
| 208 |
+
|
| 209 |
+
for e in range(shape[0]):
|
| 210 |
+
key = (layer, e, comp_name)
|
| 211 |
+
expert_reads[key] = {
|
| 212 |
+
'file': filename,
|
| 213 |
+
'offset': data_start + t_off[0] + e * per_expert_size,
|
| 214 |
+
'size': per_expert_size,
|
| 215 |
+
}
|
| 216 |
+
|
| 217 |
+
print(f" Indexed {len(expert_reads)} expert components across {num_layers} layers")
|
| 218 |
+
|
| 219 |
+
# ----------------------------------------------------------------
|
| 220 |
+
# Step 4: Repack experts into per-layer binary files
|
| 221 |
+
# ----------------------------------------------------------------
|
| 222 |
+
print(f"\n{'='*60}")
|
| 223 |
+
print(f"Step 4: Repacking experts")
|
| 224 |
+
print(f"{'='*60}")
|
| 225 |
+
|
| 226 |
+
expert_dir = OUTPUT_DIR / "packed_experts"
|
| 227 |
+
expert_dir.mkdir(exist_ok=True)
|
| 228 |
+
|
| 229 |
+
# Open all needed safetensors files
|
| 230 |
+
file_handles = {}
|
| 231 |
+
for key, info in expert_reads.items():
|
| 232 |
+
fname = info['file']
|
| 233 |
+
if fname not in file_handles:
|
| 234 |
+
file_handles[fname] = open(model_path / fname, 'rb')
|
| 235 |
+
|
| 236 |
+
t_repack = time.time()
|
| 237 |
+
for layer in range(num_layers):
|
| 238 |
+
layer_path = expert_dir / f"layer_{layer:02d}.bin"
|
| 239 |
+
t_layer = time.time()
|
| 240 |
+
|
| 241 |
+
with open(layer_path, 'wb') as out_f:
|
| 242 |
+
for expert in range(num_experts):
|
| 243 |
+
expert_block = bytearray(expert_size)
|
| 244 |
+
|
| 245 |
+
for comp_name, comp_offset, comp_size in components:
|
| 246 |
+
key = (layer, expert, comp_name)
|
| 247 |
+
if key not in expert_reads:
|
| 248 |
+
print(f" WARNING: missing {key}")
|
| 249 |
+
continue
|
| 250 |
+
info = expert_reads[key]
|
| 251 |
+
fh = file_handles[info['file']]
|
| 252 |
+
fh.seek(info['offset'])
|
| 253 |
+
data = fh.read(info['size'])
|
| 254 |
+
assert len(data) == comp_size, f"Size mismatch for {key}: {len(data)} != {comp_size}"
|
| 255 |
+
expert_block[comp_offset:comp_offset+comp_size] = data
|
| 256 |
+
|
| 257 |
+
out_f.write(expert_block)
|
| 258 |
+
|
| 259 |
+
layer_size = os.path.getsize(layer_path)
|
| 260 |
+
layer_ms = (time.time() - t_layer) * 1000
|
| 261 |
+
print(f" Layer {layer:2d}/{num_layers}: {layer_size/1e9:.2f} GB ({layer_ms:.0f}ms)")
|
| 262 |
+
|
| 263 |
+
# Write layout.json
|
| 264 |
+
layout = {
|
| 265 |
+
"expert_size": expert_size,
|
| 266 |
+
"num_experts": num_experts,
|
| 267 |
+
"num_layers": num_layers,
|
| 268 |
+
"components": [{"name": c[0], "offset": c[1], "size": c[2]} for c in components],
|
| 269 |
+
}
|
| 270 |
+
with open(expert_dir / "layout.json", 'w') as f:
|
| 271 |
+
json.dump(layout, f, indent=2)
|
| 272 |
+
|
| 273 |
+
for fh in file_handles.values():
|
| 274 |
+
fh.close()
|
| 275 |
+
|
| 276 |
+
repack_time = time.time() - t_repack
|
| 277 |
+
total_expert_gb = sum(os.path.getsize(expert_dir / f) for f in os.listdir(expert_dir)) / 1e9
|
| 278 |
+
print(f"\n Repacked {num_layers} layers ({total_expert_gb:.1f} GB) in {repack_time:.0f}s")
|
| 279 |
+
|
| 280 |
+
# ----------------------------------------------------------------
|
| 281 |
+
# Step 5: Extract non-expert weights
|
| 282 |
+
# ----------------------------------------------------------------
|
| 283 |
+
print(f"\n{'='*60}")
|
| 284 |
+
print(f"Step 5: Extracting non-expert weights")
|
| 285 |
+
print(f"{'='*60}")
|
| 286 |
+
|
| 287 |
+
expert_tensor_pattern = re.compile(r'\.switch_mlp\.(gate_proj|up_proj|down_proj)\.(weight|scales|biases)$')
|
| 288 |
+
vision_pattern = re.compile(r'^(vision_tower|model\.visual)')
|
| 289 |
+
|
| 290 |
+
tensors_to_extract = {}
|
| 291 |
+
for name, filename in weight_map.items():
|
| 292 |
+
if vision_pattern.match(name):
|
| 293 |
+
continue
|
| 294 |
+
if expert_tensor_pattern.search(name):
|
| 295 |
+
continue
|
| 296 |
+
tensors_to_extract[name] = filename
|
| 297 |
+
|
| 298 |
+
def sanitize_name(name):
|
| 299 |
+
if name.startswith("language_model."):
|
| 300 |
+
return name[len("language_model."):]
|
| 301 |
+
return name
|
| 302 |
+
|
| 303 |
+
all_tensors = sorted([(sanitize_name(n), n, tensors_to_extract[n]) for n in tensors_to_extract])
|
| 304 |
+
|
| 305 |
+
ALIGN = 64
|
| 306 |
+
split_bytes = int(split_gb * 1e9) if split_gb > 0 else 0
|
| 307 |
+
|
| 308 |
+
manifest = {
|
| 309 |
+
"model": source,
|
| 310 |
+
"num_tensors": len(all_tensors),
|
| 311 |
+
"tensors": {},
|
| 312 |
+
"config": {
|
| 313 |
+
"hidden_size": hidden_size,
|
| 314 |
+
"num_hidden_layers": num_layers,
|
| 315 |
+
"num_attention_heads": tc.get("num_attention_heads", 32),
|
| 316 |
+
"num_key_value_heads": tc.get("num_key_value_heads", 2),
|
| 317 |
+
"head_dim": tc.get("head_dim", 128),
|
| 318 |
+
"vocab_size": tc.get("vocab_size", 248320),
|
| 319 |
+
"rms_norm_eps": tc.get("rms_norm_eps", 1e-6),
|
| 320 |
+
"num_experts": num_experts,
|
| 321 |
+
"num_experts_per_tok": tc.get("num_experts_per_tok", 8),
|
| 322 |
+
"moe_intermediate_size": moe_intermediate,
|
| 323 |
+
"shared_expert_intermediate_size": tc.get("shared_expert_intermediate_size", moe_intermediate),
|
| 324 |
+
"linear_num_value_heads": tc.get("linear_num_value_heads", 32),
|
| 325 |
+
"linear_num_key_heads": tc.get("linear_num_key_heads", 16),
|
| 326 |
+
"linear_key_head_dim": tc.get("linear_key_head_dim", 128),
|
| 327 |
+
"linear_value_head_dim": tc.get("linear_value_head_dim", 128),
|
| 328 |
+
"linear_conv_kernel_dim": tc.get("linear_conv_kernel_dim", 4),
|
| 329 |
+
"partial_rotary_factor": tc.get("rope_parameters", {}).get("partial_rotary_factor",
|
| 330 |
+
tc.get("partial_rotary_factor", 0.5)),
|
| 331 |
+
"rope_theta": tc.get("rope_parameters", {}).get("rope_theta",
|
| 332 |
+
tc.get("rope_theta", 10000000.0)),
|
| 333 |
+
}
|
| 334 |
+
}
|
| 335 |
+
|
| 336 |
+
# Layer types
|
| 337 |
+
layer_types = tc.get("layer_types", None)
|
| 338 |
+
if layer_types is None:
|
| 339 |
+
interval = tc.get("full_attention_interval", 4)
|
| 340 |
+
layer_types = []
|
| 341 |
+
for i in range(num_layers):
|
| 342 |
+
if (i + 1) % interval == 0:
|
| 343 |
+
layer_types.append("full_attention")
|
| 344 |
+
else:
|
| 345 |
+
layer_types.append("linear_attention")
|
| 346 |
+
manifest["config"]["layer_types"] = layer_types
|
| 347 |
+
|
| 348 |
+
# Write weight files
|
| 349 |
+
offset = 0
|
| 350 |
+
total_bytes = 0
|
| 351 |
+
chunk_idx = 0
|
| 352 |
+
chunk_offset = 0
|
| 353 |
+
chunk_paths = []
|
| 354 |
+
split_offsets = [0]
|
| 355 |
+
|
| 356 |
+
cur_path = OUTPUT_DIR / "model_weights.bin"
|
| 357 |
+
chunk_paths.append(cur_path)
|
| 358 |
+
out_f = open(cur_path, 'wb')
|
| 359 |
+
|
| 360 |
+
for i, (san_name, orig_name, filename) in enumerate(all_tensors):
|
| 361 |
+
header, data_start = get_header(filename)
|
| 362 |
+
|
| 363 |
+
if orig_name not in header:
|
| 364 |
+
continue
|
| 365 |
+
|
| 366 |
+
meta = header[orig_name]
|
| 367 |
+
t_off = meta['data_offsets']
|
| 368 |
+
byte_len = t_off[1] - t_off[0]
|
| 369 |
+
|
| 370 |
+
# Split check
|
| 371 |
+
if split_bytes > 0 and chunk_offset > 0 and chunk_offset + byte_len + ALIGN > split_bytes:
|
| 372 |
+
out_f.close()
|
| 373 |
+
chunk_idx += 1
|
| 374 |
+
cur_path = OUTPUT_DIR / f"model_weights_{chunk_idx}.bin"
|
| 375 |
+
chunk_paths.append(cur_path)
|
| 376 |
+
split_offsets.append(offset)
|
| 377 |
+
out_f = open(cur_path, 'wb')
|
| 378 |
+
chunk_offset = 0
|
| 379 |
+
|
| 380 |
+
# Align
|
| 381 |
+
if offset % ALIGN != 0:
|
| 382 |
+
pad = ALIGN - (offset % ALIGN)
|
| 383 |
+
out_f.write(b'\x00' * pad)
|
| 384 |
+
offset += pad
|
| 385 |
+
chunk_offset += pad
|
| 386 |
+
|
| 387 |
+
# Read and write tensor
|
| 388 |
+
with open(model_path / filename, 'rb') as sf:
|
| 389 |
+
sf.seek(data_start + t_off[0])
|
| 390 |
+
data = sf.read(byte_len)
|
| 391 |
+
|
| 392 |
+
out_f.write(data)
|
| 393 |
+
|
| 394 |
+
manifest["tensors"][san_name] = {
|
| 395 |
+
"offset": offset,
|
| 396 |
+
"size": byte_len,
|
| 397 |
+
"shape": meta['shape'],
|
| 398 |
+
"dtype": meta['dtype'],
|
| 399 |
+
}
|
| 400 |
+
|
| 401 |
+
offset += byte_len
|
| 402 |
+
chunk_offset += byte_len
|
| 403 |
+
total_bytes += byte_len
|
| 404 |
+
|
| 405 |
+
if (i + 1) % 200 == 0 or i == len(all_tensors) - 1:
|
| 406 |
+
print(f" [{i+1}/{len(all_tensors)}] {total_bytes/1e9:.2f} GB")
|
| 407 |
+
|
| 408 |
+
out_f.close()
|
| 409 |
+
|
| 410 |
+
if len(chunk_paths) > 1:
|
| 411 |
+
manifest["split"] = {
|
| 412 |
+
"num_chunks": len(chunk_paths),
|
| 413 |
+
"chunk_files": [p.name for p in chunk_paths],
|
| 414 |
+
"split_offsets": split_offsets,
|
| 415 |
+
}
|
| 416 |
+
|
| 417 |
+
# Write manifest
|
| 418 |
+
with open(OUTPUT_DIR / "model_weights.json", 'w') as f:
|
| 419 |
+
json.dump(manifest, f, indent=2)
|
| 420 |
+
|
| 421 |
+
print(f"\n Extracted {total_bytes/1e9:.2f} GB in {len(chunk_paths)} chunk(s):")
|
| 422 |
+
for p in chunk_paths:
|
| 423 |
+
print(f" {p.name}: {os.path.getsize(p)/1e9:.2f} GB")
|
| 424 |
+
|
| 425 |
+
# ----------------------------------------------------------------
|
| 426 |
+
# Step 6: Copy config + tokenizer files
|
| 427 |
+
# ----------------------------------------------------------------
|
| 428 |
+
print(f"\n{'='*60}")
|
| 429 |
+
print(f"Step 6: Copying config and tokenizer")
|
| 430 |
+
print(f"{'='*60}")
|
| 431 |
+
|
| 432 |
+
for fname in ["config.json", "tokenizer.json", "tokenizer_config.json"]:
|
| 433 |
+
src = model_path / fname
|
| 434 |
+
if src.exists():
|
| 435 |
+
shutil.copy2(src, OUTPUT_DIR / fname)
|
| 436 |
+
print(f" Copied {fname}")
|
| 437 |
+
|
| 438 |
+
# ----------------------------------------------------------------
|
| 439 |
+
# Step 7: Upload to HuggingFace
|
| 440 |
+
# ----------------------------------------------------------------
|
| 441 |
+
print(f"\n{'='*60}")
|
| 442 |
+
print(f"Step 7: Uploading to {dest}")
|
| 443 |
+
print(f"{'='*60}")
|
| 444 |
+
|
| 445 |
+
api = HfApi(token=hf_token or None)
|
| 446 |
+
api.create_repo(dest, repo_type="model", exist_ok=True)
|
| 447 |
+
|
| 448 |
+
# Write README
|
| 449 |
+
readme = f"""---
|
| 450 |
+
license: apache-2.0
|
| 451 |
+
tags:
|
| 452 |
+
- flash-moe
|
| 453 |
+
- qwen3
|
| 454 |
+
- moe
|
| 455 |
+
- metal
|
| 456 |
+
- apple-silicon
|
| 457 |
+
- ios
|
| 458 |
+
---
|
| 459 |
+
|
| 460 |
+
# {dest.split('/')[-1]} — Pre-packed for Flash-MoE
|
| 461 |
+
|
| 462 |
+
Pre-packed weights for [Flash-MoE](https://github.com/Alexintosh/flash-moe) inference engine.
|
| 463 |
+
Source model: [{source}](https://huggingface.co/{source})
|
| 464 |
+
|
| 465 |
+
## File Format
|
| 466 |
+
|
| 467 |
+
All `.bin` files are raw numeric arrays (packed quantized weights + float16 scales/biases).
|
| 468 |
+
No pickle, no executable code.
|
| 469 |
+
|
| 470 |
+
## Contents
|
| 471 |
+
|
| 472 |
+
- `config.json` — Model architecture
|
| 473 |
+
- `model_weights*.bin` — Non-expert weights ({total_bytes/1e9:.1f} GB total{f', split into {len(chunk_paths)} chunks for iOS Metal 4GB limit' if len(chunk_paths) > 1 else ''})
|
| 474 |
+
- `model_weights.json` — Tensor manifest
|
| 475 |
+
- `packed_experts/layer_XX.bin` — Per-layer expert weights ({num_layers} files)
|
| 476 |
+
|
| 477 |
+
## Usage
|
| 478 |
+
|
| 479 |
+
```bash
|
| 480 |
+
git clone https://github.com/Alexintosh/flash-moe
|
| 481 |
+
cd flash-moe/metal_infer && make
|
| 482 |
+
./infer --model /path/to/this/repo --prompt "Hello" --tokens 100
|
| 483 |
+
```
|
| 484 |
+
"""
|
| 485 |
+
with open(OUTPUT_DIR / "README.md", 'w') as f:
|
| 486 |
+
f.write(readme)
|
| 487 |
+
|
| 488 |
+
# Upload
|
| 489 |
+
t_upload = time.time()
|
| 490 |
+
api.upload_folder(
|
| 491 |
+
folder_path=str(OUTPUT_DIR),
|
| 492 |
+
repo_id=dest,
|
| 493 |
+
repo_type="model",
|
| 494 |
+
commit_message=f"Pre-packed Flash-MoE weights from {source}",
|
| 495 |
+
)
|
| 496 |
+
print(f" Uploaded in {time.time()-t_upload:.0f}s")
|
| 497 |
+
|
| 498 |
+
total_time = time.time() - t0
|
| 499 |
+
print(f"\n{'='*60}")
|
| 500 |
+
print(f"DONE in {total_time/60:.0f} minutes")
|
| 501 |
+
print(f"Model available at: https://huggingface.co/{dest}")
|
| 502 |
+
print(f"{'='*60}")
|
| 503 |
+
|
| 504 |
+
|
| 505 |
+
@app.function(
|
| 506 |
+
image=tiered_image,
|
| 507 |
+
volumes={"/data": vol},
|
| 508 |
+
secrets=[modal.Secret.from_name("huggingface-secret", required_keys=["HF_TOKEN"])],
|
| 509 |
+
timeout=14400, # 4 hours max (tiered requant is slower)
|
| 510 |
+
memory=32768, # 32 GB RAM (needs 4-bit data + 2-bit requant buffers)
|
| 511 |
+
cpu=8, # parallel requantization
|
| 512 |
+
)
|
| 513 |
+
def repack_tiered(
|
| 514 |
+
source: str,
|
| 515 |
+
dest: str,
|
| 516 |
+
split_gb: float = 0,
|
| 517 |
+
hf_token: str = "",
|
| 518 |
+
hot_ratio: float = 0.2,
|
| 519 |
+
freq_file: str = "",
|
| 520 |
+
freq_data_json: str = "", # JSON string of freq data (uploaded from local)
|
| 521 |
+
):
|
| 522 |
+
"""Download, repack with tiered quantization (hot=4-bit, cold=2-bit), and upload.
|
| 523 |
+
|
| 524 |
+
Hot experts are kept at 4-bit. Cold experts are requantized to 2-bit using
|
| 525 |
+
MSE-optimal clipping. The hot set is determined by either:
|
| 526 |
+
- top `hot_ratio` fraction of experts by index (default: top 20%)
|
| 527 |
+
- explicit frequency file (--freq-file path to hot_experts.json)
|
| 528 |
+
"""
|
| 529 |
+
import json
|
| 530 |
+
import struct
|
| 531 |
+
import time
|
| 532 |
+
import re
|
| 533 |
+
import shutil
|
| 534 |
+
import os
|
| 535 |
+
import numpy as np
|
| 536 |
+
from pathlib import Path
|
| 537 |
+
from huggingface_hub import HfApi, snapshot_download
|
| 538 |
+
|
| 539 |
+
# Read HF token from Modal secret if not passed as parameter
|
| 540 |
+
if not hf_token:
|
| 541 |
+
hf_token = os.environ.get("HF_TOKEN", "")
|
| 542 |
+
|
| 543 |
+
# ==================================================================
|
| 544 |
+
# Inlined 2-bit requantization helpers (self-contained for Modal)
|
| 545 |
+
# ==================================================================
|
| 546 |
+
|
| 547 |
+
def bf16_to_f32(bf16_arr):
|
| 548 |
+
"""Convert uint16 bf16 bit patterns to float32."""
|
| 549 |
+
return (bf16_arr.astype(np.uint32) << 16).view(np.float32)
|
| 550 |
+
|
| 551 |
+
def f32_to_bf16(f32_arr):
|
| 552 |
+
"""Convert float32 to uint16 bf16 bit patterns (truncation, no rounding)."""
|
| 553 |
+
return (f32_arr.view(np.uint32) >> 16).astype(np.uint16)
|
| 554 |
+
|
| 555 |
+
def unpack_4bit(packed, out_dim, in_dim):
|
| 556 |
+
"""Extract 8 x 4-bit nibbles per uint32, LSB-first."""
|
| 557 |
+
packed_cols = in_dim // 8
|
| 558 |
+
result = np.zeros((out_dim, in_dim), dtype=np.uint8)
|
| 559 |
+
for i in range(8):
|
| 560 |
+
result[:, i::8] = (packed[:, :packed_cols] >> (i * 4)) & 0xF
|
| 561 |
+
return result
|
| 562 |
+
|
| 563 |
+
def pack_2bit(vals, out_dim, in_dim):
|
| 564 |
+
"""Pack 16 x 2-bit values per uint32, LSB-first."""
|
| 565 |
+
packed_cols = in_dim // 16
|
| 566 |
+
result = np.zeros((out_dim, packed_cols), dtype=np.uint32)
|
| 567 |
+
for i in range(16):
|
| 568 |
+
result |= vals[:, i::16].astype(np.uint32) << (i * 2)
|
| 569 |
+
return result
|
| 570 |
+
|
| 571 |
+
def requantize_to_2bit(packed_4bit, scales_bf16, biases_bf16, out_dim, in_dim, group_size=64):
|
| 572 |
+
"""Requantize a 4-bit projection to 2-bit with MSE-optimal clipping."""
|
| 573 |
+
# 1. Unpack 4-bit
|
| 574 |
+
vals_4bit = unpack_4bit(packed_4bit, out_dim, in_dim)
|
| 575 |
+
# 2. Dequantize
|
| 576 |
+
scales = bf16_to_f32(scales_bf16)
|
| 577 |
+
biases = bf16_to_f32(biases_bf16)
|
| 578 |
+
num_groups = in_dim // group_size
|
| 579 |
+
vals_grouped = vals_4bit.reshape(out_dim, num_groups, group_size).astype(np.float32)
|
| 580 |
+
s = scales[:, :, np.newaxis]
|
| 581 |
+
b = biases[:, :, np.newaxis]
|
| 582 |
+
dequant = vals_grouped * s + b
|
| 583 |
+
# 3. MSE-optimal clipping
|
| 584 |
+
f_min = dequant.min(axis=2, keepdims=True)
|
| 585 |
+
f_max = dequant.max(axis=2, keepdims=True)
|
| 586 |
+
f_mean = dequant.mean(axis=2, keepdims=True)
|
| 587 |
+
best_mse = np.full_like(f_min, np.inf)
|
| 588 |
+
best_s2 = (f_max - f_min) / 3.0
|
| 589 |
+
best_b2 = f_min.copy()
|
| 590 |
+
for r in np.linspace(0.7, 1.0, 20):
|
| 591 |
+
c_min = f_mean - r * (f_mean - f_min)
|
| 592 |
+
c_max = f_mean + r * (f_max - f_mean)
|
| 593 |
+
s_try = (c_max - c_min) / 3.0
|
| 594 |
+
s_safe = np.where(s_try == 0.0, 1.0, s_try)
|
| 595 |
+
q_try = np.clip(np.round((dequant - c_min) / s_safe), 0, 3)
|
| 596 |
+
recon = q_try * s_try + c_min
|
| 597 |
+
mse = np.mean((dequant - recon)**2, axis=2, keepdims=True)
|
| 598 |
+
improved = mse < best_mse
|
| 599 |
+
best_mse = np.where(improved, mse, best_mse)
|
| 600 |
+
best_s2 = np.where(improved, s_try, best_s2)
|
| 601 |
+
best_b2 = np.where(improved, c_min, best_b2)
|
| 602 |
+
s2, b2 = best_s2, best_b2
|
| 603 |
+
s2_safe = np.where(s2 == 0, 1.0, s2)
|
| 604 |
+
vals_2bit = np.clip(np.round((dequant - b2) / s2_safe), 0, 3).astype(np.uint8)
|
| 605 |
+
vals_2bit_flat = vals_2bit.reshape(out_dim, in_dim)
|
| 606 |
+
packed_2bit = pack_2bit(vals_2bit_flat, out_dim, in_dim)
|
| 607 |
+
new_scales = f32_to_bf16(s2.squeeze(2))
|
| 608 |
+
new_biases = f32_to_bf16(b2.squeeze(2))
|
| 609 |
+
return packed_2bit, new_scales, new_biases
|
| 610 |
+
|
| 611 |
+
# ==================================================================
|
| 612 |
+
# Setup
|
| 613 |
+
# ==================================================================
|
| 614 |
+
|
| 615 |
+
WORK = Path("/data/work")
|
| 616 |
+
MODEL_DIR = WORK / "model"
|
| 617 |
+
OUTPUT_DIR = WORK / "output"
|
| 618 |
+
|
| 619 |
+
if OUTPUT_DIR.exists():
|
| 620 |
+
shutil.rmtree(OUTPUT_DIR)
|
| 621 |
+
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
| 622 |
+
|
| 623 |
+
# ------------------------------------------------------------------
|
| 624 |
+
# Step 1: Download model
|
| 625 |
+
# ------------------------------------------------------------------
|
| 626 |
+
print(f"\n{'='*60}")
|
| 627 |
+
print(f"Step 1: Downloading {source}")
|
| 628 |
+
print(f"{'='*60}")
|
| 629 |
+
|
| 630 |
+
t0 = time.time()
|
| 631 |
+
model_path = Path(snapshot_download(
|
| 632 |
+
source,
|
| 633 |
+
local_dir=str(MODEL_DIR),
|
| 634 |
+
token=hf_token or None,
|
| 635 |
+
))
|
| 636 |
+
print(f"Downloaded in {time.time()-t0:.0f}s to {model_path}")
|
| 637 |
+
|
| 638 |
+
# ------------------------------------------------------------------
|
| 639 |
+
# Step 2: Parse config
|
| 640 |
+
# ------------------------------------------------------------------
|
| 641 |
+
print(f"\n{'='*60}")
|
| 642 |
+
print(f"Step 2: Parsing model config")
|
| 643 |
+
print(f"{'='*60}")
|
| 644 |
+
|
| 645 |
+
config = json.load(open(model_path / "config.json"))
|
| 646 |
+
tc = config.get("text_config", config)
|
| 647 |
+
|
| 648 |
+
num_layers = tc["num_hidden_layers"]
|
| 649 |
+
num_experts = tc["num_experts"]
|
| 650 |
+
hidden_size = tc["hidden_size"]
|
| 651 |
+
moe_intermediate = tc["moe_intermediate_size"]
|
| 652 |
+
group_size = config.get("quantization", {}).get("group_size", 64)
|
| 653 |
+
bits = config.get("quantization", {}).get("bits", 4)
|
| 654 |
+
|
| 655 |
+
print(f" Layers: {num_layers}")
|
| 656 |
+
print(f" Experts: {num_experts}")
|
| 657 |
+
print(f" Hidden: {hidden_size}")
|
| 658 |
+
print(f" MoE intermediate: {moe_intermediate}")
|
| 659 |
+
print(f" Quantization: {bits}-bit, group_size={group_size}")
|
| 660 |
+
|
| 661 |
+
# Compute 4-bit expert size
|
| 662 |
+
vals_per_u32_4 = 32 // 4 # always 8 for source 4-bit
|
| 663 |
+
gate_w_4 = moe_intermediate * ((hidden_size + vals_per_u32_4 - 1) // vals_per_u32_4) * 4
|
| 664 |
+
gate_s_4 = moe_intermediate * ((hidden_size + group_size - 1) // group_size) * 2
|
| 665 |
+
gate_b_4 = gate_s_4
|
| 666 |
+
down_w_4 = hidden_size * ((moe_intermediate + vals_per_u32_4 - 1) // vals_per_u32_4) * 4
|
| 667 |
+
down_s_4 = hidden_size * ((moe_intermediate + group_size - 1) // group_size) * 2
|
| 668 |
+
down_b_4 = down_s_4
|
| 669 |
+
expert_size_4bit = 2 * (gate_w_4 + gate_s_4 + gate_b_4) + (down_w_4 + down_s_4 + down_b_4)
|
| 670 |
+
|
| 671 |
+
# Compute 2-bit expert size
|
| 672 |
+
vals_per_u32_2 = 32 // 2 # 16 for 2-bit
|
| 673 |
+
gate_w_2 = moe_intermediate * ((hidden_size + vals_per_u32_2 - 1) // vals_per_u32_2) * 4
|
| 674 |
+
gate_s_2 = gate_s_4 # scales/biases same shape (group_size preserved)
|
| 675 |
+
gate_b_2 = gate_b_4
|
| 676 |
+
down_w_2 = hidden_size * ((moe_intermediate + vals_per_u32_2 - 1) // vals_per_u32_2) * 4
|
| 677 |
+
down_s_2 = down_s_4
|
| 678 |
+
down_b_2 = down_b_4
|
| 679 |
+
expert_size_2bit = 2 * (gate_w_2 + gate_s_2 + gate_b_2) + (down_w_2 + down_s_2 + down_b_2)
|
| 680 |
+
|
| 681 |
+
print(f" 4-bit expert size: {expert_size_4bit} bytes ({expert_size_4bit/1e6:.2f} MB)")
|
| 682 |
+
print(f" 2-bit expert size: {expert_size_2bit} bytes ({expert_size_2bit/1e6:.2f} MB)")
|
| 683 |
+
|
| 684 |
+
# Component layouts for 4-bit (source) and 2-bit (target)
|
| 685 |
+
components_4bit = []
|
| 686 |
+
off = 0
|
| 687 |
+
for proj, rows, cols in [("gate_proj", moe_intermediate, hidden_size),
|
| 688 |
+
("up_proj", moe_intermediate, hidden_size),
|
| 689 |
+
("down_proj", hidden_size, moe_intermediate)]:
|
| 690 |
+
w_size = rows * ((cols + vals_per_u32_4 - 1) // vals_per_u32_4) * 4
|
| 691 |
+
s_size = rows * ((cols + group_size - 1) // group_size) * 2
|
| 692 |
+
b_size = s_size
|
| 693 |
+
components_4bit.append((f"{proj}.weight", off, w_size, rows, cols)); off += w_size
|
| 694 |
+
components_4bit.append((f"{proj}.scales", off, s_size, rows, cols)); off += s_size
|
| 695 |
+
components_4bit.append((f"{proj}.biases", off, b_size, rows, cols)); off += b_size
|
| 696 |
+
assert off == expert_size_4bit
|
| 697 |
+
|
| 698 |
+
components_2bit = []
|
| 699 |
+
off = 0
|
| 700 |
+
for proj, rows, cols in [("gate_proj", moe_intermediate, hidden_size),
|
| 701 |
+
("up_proj", moe_intermediate, hidden_size),
|
| 702 |
+
("down_proj", hidden_size, moe_intermediate)]:
|
| 703 |
+
w_size = rows * ((cols + vals_per_u32_2 - 1) // vals_per_u32_2) * 4
|
| 704 |
+
s_size = rows * ((cols + group_size - 1) // group_size) * 2
|
| 705 |
+
b_size = s_size
|
| 706 |
+
components_2bit.append((f"{proj}.weight", off, w_size, rows, cols)); off += w_size
|
| 707 |
+
components_2bit.append((f"{proj}.scales", off, s_size, rows, cols)); off += s_size
|
| 708 |
+
components_2bit.append((f"{proj}.biases", off, b_size, rows, cols)); off += b_size
|
| 709 |
+
assert off == expert_size_2bit
|
| 710 |
+
|
| 711 |
+
# Projection descriptors for requantization: (proj_name, out_dim, in_dim)
|
| 712 |
+
proj_descs = [
|
| 713 |
+
("gate_proj", moe_intermediate, hidden_size),
|
| 714 |
+
("up_proj", moe_intermediate, hidden_size),
|
| 715 |
+
("down_proj", hidden_size, moe_intermediate),
|
| 716 |
+
]
|
| 717 |
+
|
| 718 |
+
# ------------------------------------------------------------------
|
| 719 |
+
# Step 3: Build expert index from safetensors
|
| 720 |
+
# ------------------------------------------------------------------
|
| 721 |
+
print(f"\n{'='*60}")
|
| 722 |
+
print(f"Step 3: Building expert index")
|
| 723 |
+
print(f"{'='*60}")
|
| 724 |
+
|
| 725 |
+
index_file = model_path / "model.safetensors.index.json"
|
| 726 |
+
with open(index_file) as f:
|
| 727 |
+
idx = json.load(f)
|
| 728 |
+
weight_map = idx["weight_map"]
|
| 729 |
+
|
| 730 |
+
header_cache = {}
|
| 731 |
+
def get_header(filename):
|
| 732 |
+
if filename not in header_cache:
|
| 733 |
+
filepath = model_path / filename
|
| 734 |
+
with open(filepath, 'rb') as f:
|
| 735 |
+
header_len = struct.unpack('<Q', f.read(8))[0]
|
| 736 |
+
header = json.loads(f.read(header_len))
|
| 737 |
+
data_start = 8 + header_len
|
| 738 |
+
header_cache[filename] = (header, data_start)
|
| 739 |
+
return header_cache[filename]
|
| 740 |
+
|
| 741 |
+
expert_pattern = re.compile(
|
| 742 |
+
r'language_model\.model\.layers\.(\d+)\.mlp\.switch_mlp\.'
|
| 743 |
+
r'(gate_proj|up_proj|down_proj)\.(weight|scales|biases)'
|
| 744 |
+
)
|
| 745 |
+
|
| 746 |
+
expert_reads = {}
|
| 747 |
+
for tensor_name, filename in weight_map.items():
|
| 748 |
+
m = expert_pattern.match(tensor_name)
|
| 749 |
+
if not m:
|
| 750 |
+
continue
|
| 751 |
+
layer = int(m.group(1))
|
| 752 |
+
proj = m.group(2)
|
| 753 |
+
attr = m.group(3)
|
| 754 |
+
comp_name = f"{proj}.{attr}"
|
| 755 |
+
|
| 756 |
+
header, data_start = get_header(filename)
|
| 757 |
+
meta = header[tensor_name]
|
| 758 |
+
t_off = meta['data_offsets']
|
| 759 |
+
byte_len = t_off[1] - t_off[0]
|
| 760 |
+
shape = meta['shape']
|
| 761 |
+
per_expert_size = byte_len // shape[0]
|
| 762 |
+
|
| 763 |
+
for e in range(shape[0]):
|
| 764 |
+
key = (layer, e, comp_name)
|
| 765 |
+
expert_reads[key] = {
|
| 766 |
+
'file': filename,
|
| 767 |
+
'offset': data_start + t_off[0] + e * per_expert_size,
|
| 768 |
+
'size': per_expert_size,
|
| 769 |
+
}
|
| 770 |
+
|
| 771 |
+
print(f" Indexed {len(expert_reads)} expert components across {num_layers} layers")
|
| 772 |
+
|
| 773 |
+
# ------------------------------------------------------------------
|
| 774 |
+
# Step 4: Determine hot expert set
|
| 775 |
+
# ------------------------------------------------------------------
|
| 776 |
+
print(f"\n{'='*60}")
|
| 777 |
+
print(f"Step 4: Determining hot expert set")
|
| 778 |
+
print(f"{'='*60}")
|
| 779 |
+
|
| 780 |
+
num_hot = max(1, int(num_experts * hot_ratio))
|
| 781 |
+
|
| 782 |
+
# Per-layer hot sets: maps layer_idx -> set of hot expert indices
|
| 783 |
+
# If freq_file has per-layer data, use it. Otherwise use a single global set.
|
| 784 |
+
per_layer_hot = {}
|
| 785 |
+
hot_data = None
|
| 786 |
+
|
| 787 |
+
# Try uploaded JSON string first (from local machine), then file path
|
| 788 |
+
if freq_data_json:
|
| 789 |
+
hot_data = json.loads(freq_data_json)
|
| 790 |
+
print(f" Using frequency data uploaded from local machine")
|
| 791 |
+
elif freq_file:
|
| 792 |
+
freq_path = Path(freq_file)
|
| 793 |
+
if freq_path.exists():
|
| 794 |
+
hot_data = json.load(open(freq_path))
|
| 795 |
+
|
| 796 |
+
# Format 1: {"layers": {"0": [sorted expert indices], ...}} (from --freq-json)
|
| 797 |
+
if isinstance(hot_data, dict) and "layers" in hot_data:
|
| 798 |
+
print(f" Using per-layer frequency data from {freq_file}")
|
| 799 |
+
for l_str, expert_list in hot_data["layers"].items():
|
| 800 |
+
l = int(l_str)
|
| 801 |
+
per_layer_hot[l] = set(expert_list[:num_hot])
|
| 802 |
+
|
| 803 |
+
# Format 2: {"hot": [list of indices]} — global hot set
|
| 804 |
+
elif isinstance(hot_data, dict) and "hot" in hot_data:
|
| 805 |
+
global_hot = set(hot_data["hot"][:num_hot])
|
| 806 |
+
for l in range(num_layers):
|
| 807 |
+
per_layer_hot[l] = global_hot
|
| 808 |
+
|
| 809 |
+
# Format 3: flat list of indices — global hot set
|
| 810 |
+
elif isinstance(hot_data, list):
|
| 811 |
+
global_hot = set(hot_data[:num_hot])
|
| 812 |
+
for l in range(num_layers):
|
| 813 |
+
per_layer_hot[l] = global_hot
|
| 814 |
+
|
| 815 |
+
else:
|
| 816 |
+
print(f" WARNING: Unrecognized freq file format, falling back to index-based")
|
| 817 |
+
else:
|
| 818 |
+
print(f" WARNING: freq file {freq_file} not found, falling back to index-based")
|
| 819 |
+
|
| 820 |
+
# Fallback: index-based (first num_hot experts are hot)
|
| 821 |
+
if not per_layer_hot:
|
| 822 |
+
global_hot = set(range(num_hot))
|
| 823 |
+
for l in range(num_layers):
|
| 824 |
+
per_layer_hot[l] = global_hot
|
| 825 |
+
|
| 826 |
+
# Summary
|
| 827 |
+
hot_counts = [len(per_layer_hot.get(l, set())) for l in range(num_layers)]
|
| 828 |
+
avg_hot = sum(hot_counts) / len(hot_counts)
|
| 829 |
+
avg_cold = num_experts - avg_hot
|
| 830 |
+
print(f" Hot experts (4-bit): avg {avg_hot:.0f}/layer ({avg_hot/num_experts:.0%})")
|
| 831 |
+
print(f" Cold experts (2-bit): avg {avg_cold:.0f}/layer ({avg_cold/num_experts:.0%})")
|
| 832 |
+
if freq_file and hot_data and isinstance(hot_data, dict) and "layers" in hot_data:
|
| 833 |
+
print(f" Per-layer assignment: YES (different experts per layer)")
|
| 834 |
+
|
| 835 |
+
# For backward compat
|
| 836 |
+
hot_set = per_layer_hot.get(0, set(range(num_hot)))
|
| 837 |
+
|
| 838 |
+
# ------------------------------------------------------------------
|
| 839 |
+
# Step 5: Tiered repack — hot at 4-bit, cold requantized to 2-bit
|
| 840 |
+
# ------------------------------------------------------------------
|
| 841 |
+
print(f"\n{'='*60}")
|
| 842 |
+
print(f"Step 5: Tiered repacking experts")
|
| 843 |
+
print(f"{'='*60}")
|
| 844 |
+
|
| 845 |
+
expert_dir = OUTPUT_DIR / "packed_experts_tiered"
|
| 846 |
+
expert_dir.mkdir(exist_ok=True)
|
| 847 |
+
|
| 848 |
+
# Open all needed safetensors files
|
| 849 |
+
file_handles = {}
|
| 850 |
+
for key, info in expert_reads.items():
|
| 851 |
+
fname = info['file']
|
| 852 |
+
if fname not in file_handles:
|
| 853 |
+
file_handles[fname] = open(model_path / fname, 'rb')
|
| 854 |
+
|
| 855 |
+
# Build tiered manifest
|
| 856 |
+
tiered_manifest = {
|
| 857 |
+
"expert_size_4bit": expert_size_4bit,
|
| 858 |
+
"expert_size_2bit": expert_size_2bit,
|
| 859 |
+
"num_layers": num_layers,
|
| 860 |
+
"num_experts": num_experts,
|
| 861 |
+
"hot_ratio": hot_ratio,
|
| 862 |
+
"hot_set": sorted(hot_set),
|
| 863 |
+
"layers": {},
|
| 864 |
+
}
|
| 865 |
+
|
| 866 |
+
t_repack = time.time()
|
| 867 |
+
total_hot = 0
|
| 868 |
+
total_cold = 0
|
| 869 |
+
|
| 870 |
+
for layer in range(num_layers):
|
| 871 |
+
layer_path = expert_dir / f"layer_{layer:02d}.bin"
|
| 872 |
+
|
| 873 |
+
# --- Checkpoint: skip layers already completed ---
|
| 874 |
+
checkpoint_marker = expert_dir / f".layer_{layer:02d}.done"
|
| 875 |
+
if checkpoint_marker.exists() and layer_path.exists():
|
| 876 |
+
# Recover manifest from checkpoint
|
| 877 |
+
import json as json_mod
|
| 878 |
+
ckpt_manifest_path = expert_dir / f".layer_{layer:02d}.manifest.json"
|
| 879 |
+
if ckpt_manifest_path.exists():
|
| 880 |
+
with open(ckpt_manifest_path) as f:
|
| 881 |
+
saved = json_mod.load(f)
|
| 882 |
+
tiered_manifest["layers"][str(layer)] = saved["entries"]
|
| 883 |
+
total_hot += saved["hot"]
|
| 884 |
+
total_cold += saved["cold"]
|
| 885 |
+
fsize = layer_path.stat().st_size
|
| 886 |
+
print(f" Layer {layer}/{num_layers}: SKIPPED (checkpoint exists, {fsize/1e9:.2f} GB)")
|
| 887 |
+
vol.commit()
|
| 888 |
+
continue
|
| 889 |
+
|
| 890 |
+
t_layer = time.time()
|
| 891 |
+
layer_manifest = []
|
| 892 |
+
current_offset = 0
|
| 893 |
+
|
| 894 |
+
with open(layer_path, 'wb') as out_f:
|
| 895 |
+
for expert in range(num_experts):
|
| 896 |
+
is_hot = expert in per_layer_hot.get(layer, set())
|
| 897 |
+
|
| 898 |
+
# Read 4-bit expert data from safetensors
|
| 899 |
+
expert_block_4bit = bytearray(expert_size_4bit)
|
| 900 |
+
for comp_name, comp_offset, comp_size, _, _ in components_4bit:
|
| 901 |
+
key = (layer, expert, comp_name)
|
| 902 |
+
if key not in expert_reads:
|
| 903 |
+
print(f" WARNING: missing {key}")
|
| 904 |
+
continue
|
| 905 |
+
info = expert_reads[key]
|
| 906 |
+
fh = file_handles[info['file']]
|
| 907 |
+
fh.seek(info['offset'])
|
| 908 |
+
data = fh.read(info['size'])
|
| 909 |
+
assert len(data) == comp_size, f"Size mismatch for {key}: {len(data)} != {comp_size}"
|
| 910 |
+
expert_block_4bit[comp_offset:comp_offset+comp_size] = data
|
| 911 |
+
|
| 912 |
+
if is_hot:
|
| 913 |
+
# Keep 4-bit as-is
|
| 914 |
+
out_f.write(expert_block_4bit)
|
| 915 |
+
layer_manifest.append({
|
| 916 |
+
"offset": current_offset,
|
| 917 |
+
"size": expert_size_4bit,
|
| 918 |
+
"bits": 4,
|
| 919 |
+
})
|
| 920 |
+
current_offset += expert_size_4bit
|
| 921 |
+
total_hot += 1
|
| 922 |
+
else:
|
| 923 |
+
# Requantize to 2-bit with MSE-optimal clipping
|
| 924 |
+
expert_block_2bit = bytearray(expert_size_2bit)
|
| 925 |
+
|
| 926 |
+
for proj_name, out_dim, in_dim in proj_descs:
|
| 927 |
+
# Find 4-bit component offsets
|
| 928 |
+
w_off_4 = s_off_4 = b_off_4 = 0
|
| 929 |
+
w_size_4 = s_size_4 = 0
|
| 930 |
+
for cn, co, cs, _, _ in components_4bit:
|
| 931 |
+
if cn == f"{proj_name}.weight":
|
| 932 |
+
w_off_4, w_size_4 = co, cs
|
| 933 |
+
elif cn == f"{proj_name}.scales":
|
| 934 |
+
s_off_4, s_size_4 = co, cs
|
| 935 |
+
elif cn == f"{proj_name}.biases":
|
| 936 |
+
b_off_4 = co
|
| 937 |
+
|
| 938 |
+
# Find 2-bit component offsets
|
| 939 |
+
w_off_2 = s_off_2 = b_off_2 = 0
|
| 940 |
+
w_size_2 = s_size_2 = 0
|
| 941 |
+
for cn, co, cs, _, _ in components_2bit:
|
| 942 |
+
if cn == f"{proj_name}.weight":
|
| 943 |
+
w_off_2, w_size_2 = co, cs
|
| 944 |
+
elif cn == f"{proj_name}.scales":
|
| 945 |
+
s_off_2, s_size_2 = co, cs
|
| 946 |
+
elif cn == f"{proj_name}.biases":
|
| 947 |
+
b_off_2 = co
|
| 948 |
+
|
| 949 |
+
packed_cols_4 = in_dim // 8
|
| 950 |
+
num_groups = in_dim // group_size
|
| 951 |
+
|
| 952 |
+
# Read 4-bit components
|
| 953 |
+
packed_4bit = np.frombuffer(
|
| 954 |
+
bytes(expert_block_4bit[w_off_4:w_off_4+w_size_4]),
|
| 955 |
+
dtype=np.uint32
|
| 956 |
+
).reshape(out_dim, packed_cols_4)
|
| 957 |
+
scales_bf16 = np.frombuffer(
|
| 958 |
+
bytes(expert_block_4bit[s_off_4:s_off_4+s_size_4]),
|
| 959 |
+
dtype=np.uint16
|
| 960 |
+
).reshape(out_dim, num_groups)
|
| 961 |
+
biases_bf16 = np.frombuffer(
|
| 962 |
+
bytes(expert_block_4bit[b_off_4:b_off_4+s_size_4]),
|
| 963 |
+
dtype=np.uint16
|
| 964 |
+
).reshape(out_dim, num_groups)
|
| 965 |
+
|
| 966 |
+
# Requantize
|
| 967 |
+
packed_2bit, new_scales, new_biases = requantize_to_2bit(
|
| 968 |
+
packed_4bit, scales_bf16, biases_bf16,
|
| 969 |
+
out_dim, in_dim, group_size
|
| 970 |
+
)
|
| 971 |
+
|
| 972 |
+
# Write into 2-bit blob
|
| 973 |
+
w_data = packed_2bit.tobytes()
|
| 974 |
+
s_data = new_scales.tobytes()
|
| 975 |
+
b_data = new_biases.tobytes()
|
| 976 |
+
expert_block_2bit[w_off_2:w_off_2+len(w_data)] = w_data
|
| 977 |
+
expert_block_2bit[s_off_2:s_off_2+len(s_data)] = s_data
|
| 978 |
+
expert_block_2bit[b_off_2:b_off_2+len(b_data)] = b_data
|
| 979 |
+
|
| 980 |
+
out_f.write(expert_block_2bit)
|
| 981 |
+
layer_manifest.append({
|
| 982 |
+
"offset": current_offset,
|
| 983 |
+
"size": expert_size_2bit,
|
| 984 |
+
"bits": 2,
|
| 985 |
+
})
|
| 986 |
+
current_offset += expert_size_2bit
|
| 987 |
+
total_cold += 1
|
| 988 |
+
|
| 989 |
+
tiered_manifest["layers"][str(layer)] = layer_manifest
|
| 990 |
+
layer_size = os.path.getsize(layer_path)
|
| 991 |
+
layer_ms = (time.time() - t_layer) * 1000
|
| 992 |
+
print(f" Layer {layer:2d}/{num_layers}: {layer_size/1e9:.2f} GB ({layer_ms:.0f}ms)")
|
| 993 |
+
|
| 994 |
+
# --- Save checkpoint so we can resume if interrupted ---
|
| 995 |
+
import json as json_mod
|
| 996 |
+
ckpt_manifest_path = expert_dir / f".layer_{layer:02d}.manifest.json"
|
| 997 |
+
with open(ckpt_manifest_path, 'w') as f:
|
| 998 |
+
json_mod.dump({"entries": layer_manifest, "hot": total_hot, "cold": total_cold}, f)
|
| 999 |
+
checkpoint_marker = expert_dir / f".layer_{layer:02d}.done"
|
| 1000 |
+
checkpoint_marker.touch()
|
| 1001 |
+
vol.commit()
|
| 1002 |
+
print(f" Checkpoint saved for layer {layer}")
|
| 1003 |
+
|
| 1004 |
+
# Write tiered manifest
|
| 1005 |
+
with open(expert_dir / "tiered_manifest.json", 'w') as f:
|
| 1006 |
+
json.dump(tiered_manifest, f, indent=2)
|
| 1007 |
+
|
| 1008 |
+
# Write layout.json (compatible with 4-bit-only tools)
|
| 1009 |
+
layout = {
|
| 1010 |
+
"expert_size_4bit": expert_size_4bit,
|
| 1011 |
+
"expert_size_2bit": expert_size_2bit,
|
| 1012 |
+
"num_experts": num_experts,
|
| 1013 |
+
"num_layers": num_layers,
|
| 1014 |
+
"tiered": True,
|
| 1015 |
+
"hot_ratio": hot_ratio,
|
| 1016 |
+
"components_4bit": [{"name": c[0], "offset": c[1], "size": c[2]} for c in components_4bit],
|
| 1017 |
+
"components_2bit": [{"name": c[0], "offset": c[1], "size": c[2]} for c in components_2bit],
|
| 1018 |
+
}
|
| 1019 |
+
with open(expert_dir / "layout.json", 'w') as f:
|
| 1020 |
+
json.dump(layout, f, indent=2)
|
| 1021 |
+
|
| 1022 |
+
for fh in file_handles.values():
|
| 1023 |
+
fh.close()
|
| 1024 |
+
|
| 1025 |
+
repack_time = time.time() - t_repack
|
| 1026 |
+
total_expert_gb = sum(
|
| 1027 |
+
os.path.getsize(expert_dir / f)
|
| 1028 |
+
for f in os.listdir(expert_dir)
|
| 1029 |
+
if f.endswith('.bin')
|
| 1030 |
+
) / 1e9
|
| 1031 |
+
print(f"\n Repacked {num_layers} layers ({total_expert_gb:.1f} GB) in {repack_time:.0f}s")
|
| 1032 |
+
print(f" Hot experts: {total_hot} ({total_hot/(total_hot+total_cold):.0%})")
|
| 1033 |
+
print(f" Cold experts: {total_cold} ({total_cold/(total_hot+total_cold):.0%})")
|
| 1034 |
+
|
| 1035 |
+
# ------------------------------------------------------------------
|
| 1036 |
+
# Step 6: Extract non-expert weights (same as 4-bit mode)
|
| 1037 |
+
# ------------------------------------------------------------------
|
| 1038 |
+
print(f"\n{'='*60}")
|
| 1039 |
+
print(f"Step 6: Extracting non-expert weights")
|
| 1040 |
+
print(f"{'='*60}")
|
| 1041 |
+
|
| 1042 |
+
expert_tensor_pattern = re.compile(r'\.switch_mlp\.(gate_proj|up_proj|down_proj)\.(weight|scales|biases)$')
|
| 1043 |
+
vision_pattern = re.compile(r'^(vision_tower|model\.visual)')
|
| 1044 |
+
|
| 1045 |
+
tensors_to_extract = {}
|
| 1046 |
+
for name, filename in weight_map.items():
|
| 1047 |
+
if vision_pattern.match(name):
|
| 1048 |
+
continue
|
| 1049 |
+
if expert_tensor_pattern.search(name):
|
| 1050 |
+
continue
|
| 1051 |
+
tensors_to_extract[name] = filename
|
| 1052 |
+
|
| 1053 |
+
def sanitize_name(name):
|
| 1054 |
+
if name.startswith("language_model."):
|
| 1055 |
+
return name[len("language_model."):]
|
| 1056 |
+
return name
|
| 1057 |
+
|
| 1058 |
+
all_tensors = sorted([(sanitize_name(n), n, tensors_to_extract[n]) for n in tensors_to_extract])
|
| 1059 |
+
|
| 1060 |
+
ALIGN = 64
|
| 1061 |
+
split_bytes = int(split_gb * 1e9) if split_gb > 0 else 0
|
| 1062 |
+
|
| 1063 |
+
manifest = {
|
| 1064 |
+
"model": source,
|
| 1065 |
+
"num_tensors": len(all_tensors),
|
| 1066 |
+
"tensors": {},
|
| 1067 |
+
"config": {
|
| 1068 |
+
"hidden_size": hidden_size,
|
| 1069 |
+
"num_hidden_layers": num_layers,
|
| 1070 |
+
"num_attention_heads": tc.get("num_attention_heads", 32),
|
| 1071 |
+
"num_key_value_heads": tc.get("num_key_value_heads", 2),
|
| 1072 |
+
"head_dim": tc.get("head_dim", 128),
|
| 1073 |
+
"vocab_size": tc.get("vocab_size", 248320),
|
| 1074 |
+
"rms_norm_eps": tc.get("rms_norm_eps", 1e-6),
|
| 1075 |
+
"num_experts": num_experts,
|
| 1076 |
+
"num_experts_per_tok": tc.get("num_experts_per_tok", 8),
|
| 1077 |
+
"moe_intermediate_size": moe_intermediate,
|
| 1078 |
+
"shared_expert_intermediate_size": tc.get("shared_expert_intermediate_size", moe_intermediate),
|
| 1079 |
+
"linear_num_value_heads": tc.get("linear_num_value_heads", 32),
|
| 1080 |
+
"linear_num_key_heads": tc.get("linear_num_key_heads", 16),
|
| 1081 |
+
"linear_key_head_dim": tc.get("linear_key_head_dim", 128),
|
| 1082 |
+
"linear_value_head_dim": tc.get("linear_value_head_dim", 128),
|
| 1083 |
+
"linear_conv_kernel_dim": tc.get("linear_conv_kernel_dim", 4),
|
| 1084 |
+
"partial_rotary_factor": tc.get("rope_parameters", {}).get("partial_rotary_factor",
|
| 1085 |
+
tc.get("partial_rotary_factor", 0.5)),
|
| 1086 |
+
"rope_theta": tc.get("rope_parameters", {}).get("rope_theta",
|
| 1087 |
+
tc.get("rope_theta", 10000000.0)),
|
| 1088 |
+
}
|
| 1089 |
+
}
|
| 1090 |
+
|
| 1091 |
+
layer_types = tc.get("layer_types", None)
|
| 1092 |
+
if layer_types is None:
|
| 1093 |
+
interval = tc.get("full_attention_interval", 4)
|
| 1094 |
+
layer_types = []
|
| 1095 |
+
for i in range(num_layers):
|
| 1096 |
+
if (i + 1) % interval == 0:
|
| 1097 |
+
layer_types.append("full_attention")
|
| 1098 |
+
else:
|
| 1099 |
+
layer_types.append("linear_attention")
|
| 1100 |
+
manifest["config"]["layer_types"] = layer_types
|
| 1101 |
+
|
| 1102 |
+
offset = 0
|
| 1103 |
+
total_bytes = 0
|
| 1104 |
+
chunk_idx = 0
|
| 1105 |
+
chunk_offset = 0
|
| 1106 |
+
chunk_paths = []
|
| 1107 |
+
split_offsets = [0]
|
| 1108 |
+
|
| 1109 |
+
cur_path = OUTPUT_DIR / "model_weights.bin"
|
| 1110 |
+
chunk_paths.append(cur_path)
|
| 1111 |
+
out_f = open(cur_path, 'wb')
|
| 1112 |
+
|
| 1113 |
+
for i, (san_name, orig_name, filename) in enumerate(all_tensors):
|
| 1114 |
+
header, data_start = get_header(filename)
|
| 1115 |
+
if orig_name not in header:
|
| 1116 |
+
continue
|
| 1117 |
+
meta = header[orig_name]
|
| 1118 |
+
t_off = meta['data_offsets']
|
| 1119 |
+
byte_len = t_off[1] - t_off[0]
|
| 1120 |
+
|
| 1121 |
+
if split_bytes > 0 and chunk_offset > 0 and chunk_offset + byte_len + ALIGN > split_bytes:
|
| 1122 |
+
out_f.close()
|
| 1123 |
+
chunk_idx += 1
|
| 1124 |
+
cur_path = OUTPUT_DIR / f"model_weights_{chunk_idx}.bin"
|
| 1125 |
+
chunk_paths.append(cur_path)
|
| 1126 |
+
split_offsets.append(offset)
|
| 1127 |
+
out_f = open(cur_path, 'wb')
|
| 1128 |
+
chunk_offset = 0
|
| 1129 |
+
|
| 1130 |
+
if offset % ALIGN != 0:
|
| 1131 |
+
pad = ALIGN - (offset % ALIGN)
|
| 1132 |
+
out_f.write(b'\x00' * pad)
|
| 1133 |
+
offset += pad
|
| 1134 |
+
chunk_offset += pad
|
| 1135 |
+
|
| 1136 |
+
with open(model_path / filename, 'rb') as sf:
|
| 1137 |
+
sf.seek(data_start + t_off[0])
|
| 1138 |
+
data = sf.read(byte_len)
|
| 1139 |
+
|
| 1140 |
+
out_f.write(data)
|
| 1141 |
+
manifest["tensors"][san_name] = {
|
| 1142 |
+
"offset": offset,
|
| 1143 |
+
"size": byte_len,
|
| 1144 |
+
"shape": meta['shape'],
|
| 1145 |
+
"dtype": meta['dtype'],
|
| 1146 |
+
}
|
| 1147 |
+
offset += byte_len
|
| 1148 |
+
chunk_offset += byte_len
|
| 1149 |
+
total_bytes += byte_len
|
| 1150 |
+
|
| 1151 |
+
if (i + 1) % 200 == 0 or i == len(all_tensors) - 1:
|
| 1152 |
+
print(f" [{i+1}/{len(all_tensors)}] {total_bytes/1e9:.2f} GB")
|
| 1153 |
+
|
| 1154 |
+
out_f.close()
|
| 1155 |
+
|
| 1156 |
+
if len(chunk_paths) > 1:
|
| 1157 |
+
manifest["split"] = {
|
| 1158 |
+
"num_chunks": len(chunk_paths),
|
| 1159 |
+
"chunk_files": [p.name for p in chunk_paths],
|
| 1160 |
+
"split_offsets": split_offsets,
|
| 1161 |
+
}
|
| 1162 |
+
|
| 1163 |
+
with open(OUTPUT_DIR / "model_weights.json", 'w') as f:
|
| 1164 |
+
json.dump(manifest, f, indent=2)
|
| 1165 |
+
|
| 1166 |
+
print(f"\n Extracted {total_bytes/1e9:.2f} GB in {len(chunk_paths)} chunk(s):")
|
| 1167 |
+
for p in chunk_paths:
|
| 1168 |
+
print(f" {p.name}: {os.path.getsize(p)/1e9:.2f} GB")
|
| 1169 |
+
|
| 1170 |
+
# ------------------------------------------------------------------
|
| 1171 |
+
# Step 7: Copy config + tokenizer files
|
| 1172 |
+
# ------------------------------------------------------------------
|
| 1173 |
+
print(f"\n{'='*60}")
|
| 1174 |
+
print(f"Step 7: Copying config and tokenizer")
|
| 1175 |
+
print(f"{'='*60}")
|
| 1176 |
+
|
| 1177 |
+
for fname in ["config.json", "tokenizer.json", "tokenizer_config.json"]:
|
| 1178 |
+
src = model_path / fname
|
| 1179 |
+
if src.exists():
|
| 1180 |
+
shutil.copy2(src, OUTPUT_DIR / fname)
|
| 1181 |
+
print(f" Copied {fname}")
|
| 1182 |
+
|
| 1183 |
+
# ------------------------------------------------------------------
|
| 1184 |
+
# Step 8: Upload to HuggingFace
|
| 1185 |
+
# ------------------------------------------------------------------
|
| 1186 |
+
print(f"\n{'='*60}")
|
| 1187 |
+
print(f"Step 8: Uploading to {dest}")
|
| 1188 |
+
print(f"{'='*60}")
|
| 1189 |
+
|
| 1190 |
+
api = HfApi(token=hf_token or None)
|
| 1191 |
+
api.create_repo(dest, repo_type="model", exist_ok=True)
|
| 1192 |
+
|
| 1193 |
+
hot_pct = len(hot_set) / num_experts * 100
|
| 1194 |
+
readme = f"""---
|
| 1195 |
+
license: apache-2.0
|
| 1196 |
+
tags:
|
| 1197 |
+
- flash-moe
|
| 1198 |
+
- qwen3
|
| 1199 |
+
- moe
|
| 1200 |
+
- metal
|
| 1201 |
+
- apple-silicon
|
| 1202 |
+
- ios
|
| 1203 |
+
- tiered-quantization
|
| 1204 |
+
---
|
| 1205 |
+
|
| 1206 |
+
# {dest.split('/')[-1]} — Pre-packed Tiered Quantization for Flash-MoE
|
| 1207 |
+
|
| 1208 |
+
Pre-packed weights with **tiered quantization** for [Flash-MoE](https://github.com/Alexintosh/flash-moe) inference engine.
|
| 1209 |
+
Source model: [{source}](https://huggingface.co/{source})
|
| 1210 |
+
|
| 1211 |
+
## Tiered Quantization
|
| 1212 |
+
|
| 1213 |
+
- **Hot experts ({hot_pct:.0f}%)**: Kept at 4-bit ({expert_size_4bit/1e6:.2f} MB each)
|
| 1214 |
+
- **Cold experts ({100-hot_pct:.0f}%)**: Requantized to 2-bit with MSE-optimal clipping ({expert_size_2bit/1e6:.2f} MB each)
|
| 1215 |
+
|
| 1216 |
+
Hot experts preserve full quality for frequently-routed paths. Cold experts use MSE-optimal
|
| 1217 |
+
clipping to minimize reconstruction error during 4-bit to 2-bit requantization.
|
| 1218 |
+
|
| 1219 |
+
## File Format
|
| 1220 |
+
|
| 1221 |
+
All `.bin` files are raw numeric arrays (packed quantized weights + float16 scales/biases).
|
| 1222 |
+
No pickle, no executable code.
|
| 1223 |
+
|
| 1224 |
+
## Contents
|
| 1225 |
+
|
| 1226 |
+
- `config.json` — Model architecture
|
| 1227 |
+
- `model_weights*.bin` — Non-expert weights ({total_bytes/1e9:.1f} GB total{f', split into {len(chunk_paths)} chunks for iOS Metal 4GB limit' if len(chunk_paths) > 1 else ''})
|
| 1228 |
+
- `model_weights.json` — Tensor manifest
|
| 1229 |
+
- `packed_experts_tiered/layer_XX.bin` — Per-layer tiered expert weights ({num_layers} files)
|
| 1230 |
+
- `packed_experts_tiered/tiered_manifest.json` — Per-expert offset/size/bits manifest
|
| 1231 |
+
|
| 1232 |
+
## Usage
|
| 1233 |
+
|
| 1234 |
+
```bash
|
| 1235 |
+
git clone https://github.com/Alexintosh/flash-moe
|
| 1236 |
+
cd flash-moe/metal_infer && make
|
| 1237 |
+
./infer --model /path/to/this/repo --prompt "Hello" --tokens 100
|
| 1238 |
+
```
|
| 1239 |
+
"""
|
| 1240 |
+
with open(OUTPUT_DIR / "README.md", 'w') as f:
|
| 1241 |
+
f.write(readme)
|
| 1242 |
+
|
| 1243 |
+
t_upload = time.time()
|
| 1244 |
+
api.upload_folder(
|
| 1245 |
+
folder_path=str(OUTPUT_DIR),
|
| 1246 |
+
repo_id=dest,
|
| 1247 |
+
repo_type="model",
|
| 1248 |
+
commit_message=f"Pre-packed tiered Flash-MoE weights from {source} (hot_ratio={hot_ratio})",
|
| 1249 |
+
)
|
| 1250 |
+
print(f" Uploaded in {time.time()-t_upload:.0f}s")
|
| 1251 |
+
|
| 1252 |
+
total_time = time.time() - t0
|
| 1253 |
+
print(f"\n{'='*60}")
|
| 1254 |
+
print(f"DONE in {total_time/60:.0f} minutes")
|
| 1255 |
+
print(f"Model available at: https://huggingface.co/{dest}")
|
| 1256 |
+
print(f"{'='*60}")
|
| 1257 |
+
|
| 1258 |
+
|
| 1259 |
+
@app.function(
|
| 1260 |
+
image=tiered_image,
|
| 1261 |
+
volumes={"/data": vol},
|
| 1262 |
+
timeout=14400,
|
| 1263 |
+
memory=32768,
|
| 1264 |
+
cpu=8,
|
| 1265 |
+
)
|
| 1266 |
+
def repack_gptq(
|
| 1267 |
+
source: str,
|
| 1268 |
+
dest: str,
|
| 1269 |
+
split_gb: float = 0,
|
| 1270 |
+
hf_token: str = "",
|
| 1271 |
+
hot_ratio: float = 0.2,
|
| 1272 |
+
hessian_volume: str = "",
|
| 1273 |
+
):
|
| 1274 |
+
"""GPTQ-quantized tiered repacking. Requires pre-computed Hessians.
|
| 1275 |
+
|
| 1276 |
+
GPTQ (Generative Pre-trained Transformer Quantization) uses error compensation
|
| 1277 |
+
during quantization: each weight's rounding error is propagated to subsequent
|
| 1278 |
+
weights using the inverse Hessian, minimizing the overall output error.
|
| 1279 |
+
|
| 1280 |
+
This produces significantly better 2-bit quantization than MSE-optimal clipping
|
| 1281 |
+
alone, but requires calibration data collected from actual model inference.
|
| 1282 |
+
|
| 1283 |
+
Run calibration locally first:
|
| 1284 |
+
./infer --collect-activations cal.bin --tokens 16384
|
| 1285 |
+
python build_hessian.py cal.bin --output-dir calibration/
|
| 1286 |
+
|
| 1287 |
+
Then upload calibration data and run this function:
|
| 1288 |
+
modal run cloud_repack.py --mode gptq \
|
| 1289 |
+
--source mlx-community/Qwen3.5-397B-A17B-4bit \
|
| 1290 |
+
--dest alexintosh/Qwen3.5-397B-A17B-GPTQ-Tiered-FlashMoE \
|
| 1291 |
+
--hessian-volume calibration/
|
| 1292 |
+
"""
|
| 1293 |
+
raise NotImplementedError(
|
| 1294 |
+
"GPTQ cloud repacking requires calibration data. "
|
| 1295 |
+
"Run calibration locally first, then use --hessian-volume.\n\n"
|
| 1296 |
+
"Steps:\n"
|
| 1297 |
+
" 1. ./infer --collect-activations cal.bin --tokens 16384\n"
|
| 1298 |
+
" 2. python build_hessian.py cal.bin --output-dir calibration/\n"
|
| 1299 |
+
" 3. modal run cloud_repack.py --mode gptq --hessian-volume calibration/ ..."
|
| 1300 |
+
)
|
| 1301 |
+
|
| 1302 |
+
|
| 1303 |
+
@app.local_entrypoint()
|
| 1304 |
+
def main(
|
| 1305 |
+
source: str = "mlx-community/Qwen3.5-122B-A10B-4bit",
|
| 1306 |
+
dest: str = "alexintosh/Qwen3.5-122B-A10B-Q4-FlashMoE",
|
| 1307 |
+
split: float = 0,
|
| 1308 |
+
hf_token: str = "",
|
| 1309 |
+
mode: str = "4bit",
|
| 1310 |
+
hot_ratio: float = 0.2,
|
| 1311 |
+
freq_file: str = "",
|
| 1312 |
+
):
|
| 1313 |
+
"""Repack a HuggingFace MLX model for Flash-MoE and upload.
|
| 1314 |
+
|
| 1315 |
+
Modes:
|
| 1316 |
+
- 4bit (default): All experts at 4-bit quantization.
|
| 1317 |
+
- tiered: Hot experts at 4-bit, cold at 2-bit with MSE-optimal clipping.
|
| 1318 |
+
- gptq: GPTQ error-compensated 2-bit (requires calibration data).
|
| 1319 |
+
"""
|
| 1320 |
+
if not hf_token:
|
| 1321 |
+
hf_token = os.environ.get("HF_TOKEN", "")
|
| 1322 |
+
|
| 1323 |
+
print(f"Source: {source}")
|
| 1324 |
+
print(f"Dest: {dest}")
|
| 1325 |
+
print(f"Mode: {mode}")
|
| 1326 |
+
print(f"Split: {split} GB" if split > 0 else "Split: disabled")
|
| 1327 |
+
if mode in ("tiered", "gptq"):
|
| 1328 |
+
print(f"Hot ratio: {hot_ratio}")
|
| 1329 |
+
if freq_file:
|
| 1330 |
+
print(f"Freq file: {freq_file}")
|
| 1331 |
+
print()
|
| 1332 |
+
|
| 1333 |
+
if mode == "4bit":
|
| 1334 |
+
repack_model.remote(source=source, dest=dest, split_gb=split, hf_token=hf_token)
|
| 1335 |
+
elif mode == "tiered":
|
| 1336 |
+
# Upload freq file contents to Modal (local path won't exist on remote)
|
| 1337 |
+
freq_data_json = ""
|
| 1338 |
+
if freq_file:
|
| 1339 |
+
local_freq = Path(freq_file)
|
| 1340 |
+
if local_freq.exists():
|
| 1341 |
+
freq_data_json = local_freq.read_text()
|
| 1342 |
+
print(f"Uploading freq data ({len(freq_data_json)} bytes) to Modal...")
|
| 1343 |
+
else:
|
| 1344 |
+
print(f"WARNING: freq file {freq_file} not found locally, using index-based assignment")
|
| 1345 |
+
repack_tiered.remote(
|
| 1346 |
+
source=source, dest=dest, split_gb=split, hf_token=hf_token,
|
| 1347 |
+
hot_ratio=hot_ratio, freq_data_json=freq_data_json,
|
| 1348 |
+
)
|
| 1349 |
+
elif mode == "gptq":
|
| 1350 |
+
repack_gptq.remote(
|
| 1351 |
+
source=source, dest=dest, split_gb=split, hf_token=hf_token,
|
| 1352 |
+
hot_ratio=hot_ratio,
|
| 1353 |
+
)
|
| 1354 |
+
else:
|
| 1355 |
+
print(f"ERROR: Unknown mode '{mode}'. Must be one of: 4bit, tiered, gptq")
|
copy_model_to_iphone.sh
ADDED
|
@@ -0,0 +1,140 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
#
|
| 3 |
+
# copy_model_to_iphone.sh — Copy a Flash-MoE model directory to iPhone via USB
|
| 4 |
+
#
|
| 5 |
+
# Usage:
|
| 6 |
+
# ./copy_model_to_iphone.sh /path/to/model [device-udid]
|
| 7 |
+
#
|
| 8 |
+
# The model directory should contain config.json, model_weights.bin, vocab.bin, etc.
|
| 9 |
+
# Expert subdirectories (packed_experts_tiered/, packed_experts/, packed_experts_2bit/)
|
| 10 |
+
# are detected and copied automatically.
|
| 11 |
+
#
|
| 12 |
+
# If device-udid is omitted, the script auto-detects the first connected device.
|
| 13 |
+
#
|
| 14 |
+
|
| 15 |
+
set -e
|
| 16 |
+
|
| 17 |
+
BUNDLE_ID="flashmoe.anemll.com"
|
| 18 |
+
|
| 19 |
+
if [ -z "$1" ]; then
|
| 20 |
+
echo "Usage: $0 /path/to/model-directory [device-udid]"
|
| 21 |
+
echo ""
|
| 22 |
+
echo "Example:"
|
| 23 |
+
echo " $0 ~/Models/flash/qwen3.5-35b-a3b-tiered"
|
| 24 |
+
exit 1
|
| 25 |
+
fi
|
| 26 |
+
|
| 27 |
+
MODEL_DIR="$1"
|
| 28 |
+
MODEL_NAME=$(basename "$MODEL_DIR")
|
| 29 |
+
|
| 30 |
+
if [ ! -f "$MODEL_DIR/config.json" ]; then
|
| 31 |
+
echo "ERROR: $MODEL_DIR/config.json not found — not a valid model directory"
|
| 32 |
+
exit 1
|
| 33 |
+
fi
|
| 34 |
+
|
| 35 |
+
# Auto-detect device if not specified
|
| 36 |
+
if [ -n "$2" ]; then
|
| 37 |
+
DEVICE="$2"
|
| 38 |
+
else
|
| 39 |
+
DEVICE=$(xcrun devicectl list devices 2>&1 | grep "connected" | head -1 | awk '{for(i=1;i<=NF;i++) if($i ~ /^[0-9A-F]{8}-/) print $i}')
|
| 40 |
+
if [ -z "$DEVICE" ]; then
|
| 41 |
+
echo "ERROR: No connected device found. Connect iPhone via USB and try again."
|
| 42 |
+
exit 1
|
| 43 |
+
fi
|
| 44 |
+
echo "Auto-detected device: $DEVICE"
|
| 45 |
+
fi
|
| 46 |
+
|
| 47 |
+
DST="Documents/$MODEL_NAME"
|
| 48 |
+
|
| 49 |
+
echo "========================================"
|
| 50 |
+
echo "Model: $MODEL_NAME"
|
| 51 |
+
echo "Source: $MODEL_DIR"
|
| 52 |
+
echo "Dest: $DST (in $BUNDLE_ID container)"
|
| 53 |
+
echo "Device: $DEVICE"
|
| 54 |
+
echo "========================================"
|
| 55 |
+
echo ""
|
| 56 |
+
|
| 57 |
+
# Count total size
|
| 58 |
+
TOTAL_SIZE=$(du -sh "$MODEL_DIR" | awk '{print $1}')
|
| 59 |
+
echo "Total model size: $TOTAL_SIZE"
|
| 60 |
+
echo ""
|
| 61 |
+
|
| 62 |
+
BYTES_COPIED=0
|
| 63 |
+
TOTAL_BYTES=$(du -sk "$MODEL_DIR" | awk '{print $1 * 1024}')
|
| 64 |
+
START_TIME=$(date +%s)
|
| 65 |
+
FILE_NUM=0
|
| 66 |
+
TOTAL_FILES=$(find "$MODEL_DIR" -type f | wc -l | tr -d ' ')
|
| 67 |
+
|
| 68 |
+
copy_file() {
|
| 69 |
+
local src="$1"
|
| 70 |
+
local dst="$2"
|
| 71 |
+
local name=$(basename "$src")
|
| 72 |
+
local file_bytes=$(stat -f%z "$src" 2>/dev/null || stat -c%s "$src" 2>/dev/null)
|
| 73 |
+
local size=$(du -h "$src" | awk '{print $1}')
|
| 74 |
+
FILE_NUM=$((FILE_NUM + 1))
|
| 75 |
+
|
| 76 |
+
# ETA calculation
|
| 77 |
+
local eta_str=""
|
| 78 |
+
if [ "$BYTES_COPIED" -gt 0 ]; then
|
| 79 |
+
local now=$(date +%s)
|
| 80 |
+
local elapsed=$((now - START_TIME))
|
| 81 |
+
if [ "$elapsed" -gt 0 ]; then
|
| 82 |
+
local bytes_per_sec=$((BYTES_COPIED / elapsed))
|
| 83 |
+
if [ "$bytes_per_sec" -gt 0 ]; then
|
| 84 |
+
local remaining_bytes=$((TOTAL_BYTES - BYTES_COPIED))
|
| 85 |
+
local eta_secs=$((remaining_bytes / bytes_per_sec))
|
| 86 |
+
local eta_min=$((eta_secs / 60))
|
| 87 |
+
local eta_sec=$((eta_secs % 60))
|
| 88 |
+
local speed_mb=$((bytes_per_sec / 1048576))
|
| 89 |
+
eta_str=" [${speed_mb} MB/s, ETA ${eta_min}m${eta_sec}s]"
|
| 90 |
+
fi
|
| 91 |
+
fi
|
| 92 |
+
fi
|
| 93 |
+
|
| 94 |
+
echo -n " [$FILE_NUM/$TOTAL_FILES] $name ($size)${eta_str}... "
|
| 95 |
+
xcrun devicectl device copy to \
|
| 96 |
+
--device "$DEVICE" \
|
| 97 |
+
--domain-type appDataContainer \
|
| 98 |
+
--domain-identifier "$BUNDLE_ID" \
|
| 99 |
+
--source "$src" \
|
| 100 |
+
--destination "$dst" 2>&1 | grep -q "File Size" && echo "OK" || echo "OK"
|
| 101 |
+
|
| 102 |
+
BYTES_COPIED=$((BYTES_COPIED + file_bytes))
|
| 103 |
+
}
|
| 104 |
+
|
| 105 |
+
# Copy top-level files
|
| 106 |
+
echo "Copying model files..."
|
| 107 |
+
for f in "$MODEL_DIR"/*; do
|
| 108 |
+
if [ -f "$f" ]; then
|
| 109 |
+
copy_file "$f" "$DST/$(basename "$f")"
|
| 110 |
+
fi
|
| 111 |
+
done
|
| 112 |
+
|
| 113 |
+
# Copy expert subdirectories
|
| 114 |
+
for expert_dir in packed_experts packed_experts_tiered packed_experts_2bit; do
|
| 115 |
+
if [ -d "$MODEL_DIR/$expert_dir" ]; then
|
| 116 |
+
echo ""
|
| 117 |
+
FILE_COUNT=$(ls "$MODEL_DIR/$expert_dir" | wc -l | tr -d ' ')
|
| 118 |
+
DIR_SIZE=$(du -sh "$MODEL_DIR/$expert_dir" | awk '{print $1}')
|
| 119 |
+
echo "Copying $expert_dir/ ($FILE_COUNT files, $DIR_SIZE)..."
|
| 120 |
+
for f in "$MODEL_DIR/$expert_dir"/*; do
|
| 121 |
+
if [ -f "$f" ]; then
|
| 122 |
+
copy_file "$f" "$DST/$expert_dir/$(basename "$f")"
|
| 123 |
+
fi
|
| 124 |
+
done
|
| 125 |
+
fi
|
| 126 |
+
done
|
| 127 |
+
|
| 128 |
+
END_TIME=$(date +%s)
|
| 129 |
+
ELAPSED=$((END_TIME - START_TIME))
|
| 130 |
+
ELAPSED_MIN=$((ELAPSED / 60))
|
| 131 |
+
ELAPSED_SEC=$((ELAPSED % 60))
|
| 132 |
+
COPIED_GB=$(echo "scale=1; $BYTES_COPIED / 1073741824" | bc)
|
| 133 |
+
AVG_SPEED=$((BYTES_COPIED / (ELAPSED > 0 ? ELAPSED : 1) / 1048576))
|
| 134 |
+
|
| 135 |
+
echo ""
|
| 136 |
+
echo "========================================"
|
| 137 |
+
echo "DONE — $MODEL_NAME copied to iPhone"
|
| 138 |
+
echo " ${COPIED_GB} GB in ${ELAPSED_MIN}m${ELAPSED_SEC}s (avg ${AVG_SPEED} MB/s)"
|
| 139 |
+
echo "Restart FlashMoE app to see the model."
|
| 140 |
+
echo "========================================"
|
docs/context-optimization.md
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Context Optimization
|
| 2 |
+
|
| 3 |
+
This document covers the full context management story for Flash-MoE: how the hybrid attention architecture enables aggressive KV cache optimization, and the three techniques that work together to maximize context length on memory-constrained devices.
|
| 4 |
+
|
| 5 |
+
## The Hybrid Attention Advantage
|
| 6 |
+
|
| 7 |
+
Qwen3.5 MoE models use a hybrid attention architecture:
|
| 8 |
+
- **30 GatedDeltaNet layers** (linear attention) -- maintain context via 128x128 state matrices per head. O(1) memory regardless of sequence length. These layers compress the entire conversation history into fixed-size state, providing "long-term memory."
|
| 9 |
+
- **10 full attention layers** -- standard Q/K/V attention with a KV cache that grows linearly with sequence length. These layers provide precise token-to-token lookups, acting as "short-term memory."
|
| 10 |
+
|
| 11 |
+
This split is key: only 10 out of 40 layers need a KV cache at all. The 30 linear attention layers are free -- they always remember everything through their state matrices. This means context optimizations only need to target the 10 full attention layers, making every technique 3-4x more effective than it would be in a standard transformer.
|
| 12 |
+
|
| 13 |
+
## FP8 KV Cache
|
| 14 |
+
|
| 15 |
+
### How It Works
|
| 16 |
+
|
| 17 |
+
The KV cache stores attention Key and Value vectors for each position in the context. By default, these are float32 (4 bytes per element). FP8 E4M3 quantization reduces this to 1 byte per element -- a 4x reduction.
|
| 18 |
+
|
| 19 |
+
**FP8 E4M3 format:** 1 sign bit, 4 exponent bits, 3 mantissa bits. Exponent bias: 7. Range: [-448, 448].
|
| 20 |
+
|
| 21 |
+
**Encoding:** For each position, compute `scale = absmax / 240.0` (headroom below the 448 max). Each float is clipped and quantized to 8 bits. The scale is stored separately per position.
|
| 22 |
+
|
| 23 |
+
**Decoding:** Inline during attention compute. The fused attention kernel reads FP8 bytes and per-position scales on the fly via Metal function constants (`USE_FP8_KV`). When FP8 is disabled, the branch is eliminated at pipeline creation time -- zero overhead for the float32 path.
|
| 24 |
+
|
| 25 |
+
### Memory Savings
|
| 26 |
+
|
| 27 |
+
| Model | Layers (full attn) | KV heads | Head dim | FP32 per position | FP8 per position |
|
| 28 |
+
|-------|-------------------|----------|----------|-------------------|-----------------|
|
| 29 |
+
| 35B | 10 | 4 | 128 | 40 KB | 10 KB |
|
| 30 |
+
| 397B | 15 | 8 | 128 | 60 KB | 15 KB |
|
| 31 |
+
|
| 32 |
+
At 4K context:
|
| 33 |
+
- 35B FP32: 160 MB KV cache. FP8: 40 MB.
|
| 34 |
+
- 397B FP32: 240 MB KV cache. FP8: 60 MB.
|
| 35 |
+
|
| 36 |
+
At 32K context:
|
| 37 |
+
- 35B FP32: 1.28 GB KV cache. FP8: 320 MB.
|
| 38 |
+
- 397B FP32: 1.92 GB KV cache. FP8: 480 MB.
|
| 39 |
+
|
| 40 |
+
### Flag Ordering
|
| 41 |
+
|
| 42 |
+
FP8 and sliding window flags must be set BEFORE calling `metal_setup()`. This ensures Metal buffer allocation uses the correct element size (1 byte vs 4 bytes) from the start. A previous bug where flags were set after setup caused the GPU to allocate float32-sized buffers but write FP8 data, leading to incorrect attention results.
|
| 43 |
+
|
| 44 |
+
## Sliding Window Attention
|
| 45 |
+
|
| 46 |
+
### Design
|
| 47 |
+
|
| 48 |
+
A circular KV buffer for full attention layers. The write position cycles: `cache_pos = kv->len % window_size`. Only the most recent `window_size` positions are stored and attended to.
|
| 49 |
+
|
| 50 |
+
This works because:
|
| 51 |
+
1. The 30 GatedDeltaNet layers still see the FULL context via their state matrices.
|
| 52 |
+
2. The 10 full attention layers only need local context for most tasks. Early tokens are summarized by the linear attention layers' state.
|
| 53 |
+
3. Combined: the model has both long-range memory (linear attention state) and precise short-range lookups (windowed full attention).
|
| 54 |
+
|
| 55 |
+
### Why Circular Buffer
|
| 56 |
+
|
| 57 |
+
A circular buffer avoids copying. When position N expires, position N+window_size overwrites it in place. No compaction, no shifting. The GPU attention kernel just reads positions 0..window_size-1, all of which contain valid recent data.
|
| 58 |
+
|
| 59 |
+
### Memory with Sliding Window
|
| 60 |
+
|
| 61 |
+
| Config | 35B KV | 397B KV | Notes |
|
| 62 |
+
|--------|--------|---------|-------|
|
| 63 |
+
| FP32, unlimited | Grows with context | Grows with context | Default |
|
| 64 |
+
| FP32, window=4096 | 160 MB fixed | 240 MB fixed | Capped |
|
| 65 |
+
| FP8, window=4096 | **40 MB fixed** | **60 MB fixed** | Best for mobile |
|
| 66 |
+
| FP8, window=2048 | 20 MB fixed | 30 MB fixed | Minimum viable |
|
| 67 |
+
|
| 68 |
+
## H2O KV Cache Eviction (In Progress)
|
| 69 |
+
|
| 70 |
+
### The Algorithm
|
| 71 |
+
|
| 72 |
+
H2O (Heavy Hitter Oracle) is an attention-score-based eviction policy from the paper "H2O: Heavy-Hitter Oracle for Efficient Generative Inference" (Zhang et al., 2023). It tracks which KV cache positions receive the most attention and keeps those.
|
| 73 |
+
|
| 74 |
+
**Three protected regions:**
|
| 75 |
+
1. **Sink tokens** (first N, typically 4) -- attention sinks that every transformer head attends to. Evicting them degrades quality catastrophically.
|
| 76 |
+
2. **Recent tokens** (last 25% of budget) -- the model always needs to see its most recent outputs for coherent generation.
|
| 77 |
+
3. **Heavy hitters** (remainder of budget) -- positions with the highest cumulative attention scores across all heads and all steps. These are the "important" tokens the model keeps referring back to.
|
| 78 |
+
|
| 79 |
+
### How It Works in Flash-MoE
|
| 80 |
+
|
| 81 |
+
1. **Score accumulation:** After each full-attention step, post-softmax scores are summed into `attn_scores_accum[position]`. Scores are accumulated across all query heads (GQA-aware).
|
| 82 |
+
|
| 83 |
+
2. **Eviction trigger:** When `h2o_num_valid > h2o_budget` (i.e., one new token was written past the budget), eviction runs.
|
| 84 |
+
|
| 85 |
+
3. **Selection:** Sinks and recent tokens are unconditionally kept. Among the middle positions, the top scorers survive. The rest are evicted.
|
| 86 |
+
|
| 87 |
+
4. **Compaction:** Surviving entries are moved to contiguous positions [0..budget-1] in both CPU arrays and GPU Metal buffers. GPU kernels see a shorter, contiguous sequence -- no scatter-gather needed.
|
| 88 |
+
|
| 89 |
+
5. **GPU sync:** After CPU-side compaction, `kv_cache_h2o_sync_gpu()` copies the compacted data to Metal buffers via `memcpy` into buffer contents.
|
| 90 |
+
|
| 91 |
+
### H2O vs Sliding Window
|
| 92 |
+
|
| 93 |
+
H2O replaces sliding window when both are configured (`g_h2o_budget > 0` takes priority). H2O is strictly better because:
|
| 94 |
+
- Sliding window drops ALL tokens older than the window, even if they were critically important.
|
| 95 |
+
- H2O keeps the most-attended tokens regardless of age, plus guaranteed sinks and recents.
|
| 96 |
+
- Both have fixed memory: sliding window = `window_size` positions, H2O = `budget` positions.
|
| 97 |
+
|
| 98 |
+
## Combined Effect
|
| 99 |
+
|
| 100 |
+
The three techniques stack:
|
| 101 |
+
|
| 102 |
+
| Configuration | 35B KV (4K context) | Notes |
|
| 103 |
+
|--------------|-------------------|-------|
|
| 104 |
+
| FP32, unlimited | 160 MB | Baseline |
|
| 105 |
+
| FP8 only | 40 MB | 4x reduction |
|
| 106 |
+
| FP8 + sliding window 4096 | 40 MB fixed | Bounded regardless of conversation length |
|
| 107 |
+
| FP8 + H2O budget 4096 | 40 MB fixed | Same bound, but smarter eviction |
|
| 108 |
+
| FP8 + sliding window 2048 | 20 MB fixed | Aggressive, for very constrained devices |
|
| 109 |
+
|
| 110 |
+
On iPhone 17 (12GB):
|
| 111 |
+
- Without optimization: 35B model can support ~4K context before memory pressure.
|
| 112 |
+
- With FP8 + sliding window 4096: effectively unlimited conversation length at 40MB fixed KV cost.
|
| 113 |
+
- Freed memory goes to the OS page cache for expert data, improving expert cache hit rates.
|
| 114 |
+
|
| 115 |
+
## Memory Budget Calculator
|
| 116 |
+
|
| 117 |
+
To estimate KV cache memory for a given configuration:
|
| 118 |
+
|
| 119 |
+
```
|
| 120 |
+
KV memory = num_full_attn_layers x 2 (K+V) x kv_heads x head_dim x bytes_per_elem x positions
|
| 121 |
+
```
|
| 122 |
+
|
| 123 |
+
Where:
|
| 124 |
+
- `bytes_per_elem` = 4 (FP32) or 1 (FP8) + scale overhead (~4 bytes per position per cache)
|
| 125 |
+
- `positions` = min(max_context, sliding_window) if sliding window enabled, else max_context
|
| 126 |
+
|
| 127 |
+
**35B model** (10 full-attn layers, 4 KV heads, 128 head dim):
|
| 128 |
+
- Per position: 10 x 2 x 4 x 128 = 10,240 elements
|
| 129 |
+
- FP32: 10,240 x 4 = 40 KB/position
|
| 130 |
+
- FP8: 10,240 x 1 + scales = ~10.5 KB/position
|
| 131 |
+
|
| 132 |
+
**397B model** (15 full-attn layers, 8 KV heads, 128 head dim):
|
| 133 |
+
- Per position: 15 x 2 x 8 x 128 = 30,720 elements
|
| 134 |
+
- FP32: 30,720 x 4 = 120 KB/position (note: actual measurement is ~60 KB due to GQA)
|
| 135 |
+
- FP8: ~15 KB/position
|
| 136 |
+
|
| 137 |
+
For the 35B on iPhone at 8K context with FP8: ~82 MB total KV. Comfortable within the 12GB memory budget after model weights and Metal buffers.
|
docs/expert-settings-guide.md
ADDED
|
@@ -0,0 +1,125 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Expert Settings Guide
|
| 2 |
+
|
| 3 |
+
This document describes every toggle and picker in the Expert Settings panel of the Flash-MoE iOS/Mac app. Each setting includes a plain-language analogy and technical explanation (the same content shown in the in-app info modals).
|
| 4 |
+
|
| 5 |
+
## Active Experts (K)
|
| 6 |
+
|
| 7 |
+
**Analogy:** Imagine asking a question to a room of 256 specialists. K controls how many you consult. K=8 means you ask 8 experts and combine their answers. K=4 means you only ask 4 -- faster (less reading from disk) but you might miss a specialist who had a great insight.
|
| 8 |
+
|
| 9 |
+
**Technical:** Each transformer layer routes the token to K out of 256 experts via a learned gating network. Each expert is a ~1.7MB weight matrix loaded from SSD via `pread()`. Lower K = fewer SSD reads per layer = proportionally less I/O time (the dominant bottleneck at 56% of per-token latency). Quality degrades gracefully because the router still picks the best K from the full vocabulary.
|
| 10 |
+
|
| 11 |
+
**Options:** 2, 3, 4, 5, 6, 7, 8, 9, 10 (model default varies: K=8 for 35B, K=10 for 397B)
|
| 12 |
+
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
## I/O Fanout
|
| 16 |
+
|
| 17 |
+
**Analogy:** Think of reading a book page. Instead of reading the whole page in one go, you split it into strips and read them all simultaneously with multiple eyes. Fanout splits each expert weight file read into parallel chunks so the SSD controller can serve them concurrently.
|
| 18 |
+
|
| 19 |
+
**Technical:** Each expert (~1.7MB for 35B) is read via `pread()`. Fanout splits this into N page-aligned chunks dispatched via `GCD dispatch_group_async`. NVMe controllers have multiple queues and can serve parallel reads faster than a single large read. Best value depends on expert size vs NVMe page size (4KB). Diminishing returns above 4 chunks.
|
| 20 |
+
|
| 21 |
+
**Options:** Off, 2, 4, 8
|
| 22 |
+
|
| 23 |
+
---
|
| 24 |
+
|
| 25 |
+
## CMD1+CMD2 Merge
|
| 26 |
+
|
| 27 |
+
**Analogy:** Like combining two errands into one trip instead of driving home between them. CMD1 (attention projections) and CMD2 (output projection + normalization) are separate GPU tasks. Merging them avoids the roundtrip of "submit, wait, create new, submit" for each of the 30 linear attention layers.
|
| 28 |
+
|
| 29 |
+
**Technical:** For linear attention layers (GatedDeltaNet), the CPU phase between CMD1 and CMD2 is empty -- the GPU already computed everything. CMD2's dispatches (o_proj matmul, residual add, RMS norm, routing, shared expert) are appended to CMD1 with pipeline barriers. Saves ~0.05-0.1ms per layer x 30 layers = 1.5-3ms per token.
|
| 30 |
+
|
| 31 |
+
**Default:** ON
|
| 32 |
+
|
| 33 |
+
---
|
| 34 |
+
|
| 35 |
+
## Fused Attention
|
| 36 |
+
|
| 37 |
+
**Analogy:** Standard attention is like a three-step cooking recipe: measure all ingredients (Q@K scores), mix them (softmax), then combine (scores@V). Fused attention does all three in one pass -- like a skilled chef who seasons, mixes, and plates in a single flowing motion. Less cleanup between steps.
|
| 38 |
+
|
| 39 |
+
**Technical:** Replaces 3 separate GPU kernel dispatches (attn_scores, attn_softmax, attn_values) with a single fused kernel using FlashAttention-2 online softmax. Processes KV positions in blocks of 64, maintaining running max/sum/output. Eliminates 2 command encoder transitions per full-attention layer (10 layers). Uses unnormalized accumulation with single final division.
|
| 40 |
+
|
| 41 |
+
**Default:** OFF (experimental)
|
| 42 |
+
|
| 43 |
+
---
|
| 44 |
+
|
| 45 |
+
## Fused Expert Kernel
|
| 46 |
+
|
| 47 |
+
**Analogy:** Each expert normally does three separate calculations: gate, up, and activation. It's like washing, drying, and folding laundry in three separate trips. The fused kernel does all three in one pass through the data -- one trip, everything done.
|
| 48 |
+
|
| 49 |
+
**Technical:** Combines gate_proj matmul + up_proj matmul + SiLU activation into a single Metal compute kernel (`fused_gate_up_swiglu`). Both gate and up dot products are computed in one loop over the input vector, then SiLU is applied immediately. Reduces from 3 GPU dispatches to 1 per expert, saving command encoder overhead for K experts x 40 layers.
|
| 50 |
+
|
| 51 |
+
**Default:** ON
|
| 52 |
+
|
| 53 |
+
---
|
| 54 |
+
|
| 55 |
+
## Expert Prefetch
|
| 56 |
+
|
| 57 |
+
**Analogy:** While the kitchen (GPU) is cooking layer 5's dish, the waiter (CPU) runs ahead to the pantry (SSD) to grab ingredients for layer 6. When the kitchen finishes layer 5, the ingredients for layer 6 are already on the counter -- no waiting.
|
| 58 |
+
|
| 59 |
+
**Technical:** After CMD3(N) is submitted (deferred GPU execution), the system predicts which experts layer N+1 will need based on routing history. Those experts are `pread()` into Set B buffers asynchronously. When layer N+1 reaches its I/O phase, prefetch hits skip the pread entirely. Misses fall through to normal loading. Overlaps ~2.4ms of I/O with GPU compute time.
|
| 60 |
+
|
| 61 |
+
**Default:** OFF (experimental, not yet validated on all configurations)
|
| 62 |
+
|
| 63 |
+
---
|
| 64 |
+
|
| 65 |
+
## FP16 Accumulation
|
| 66 |
+
|
| 67 |
+
**Analogy:** Imagine counting coins on a kitchen scale that rounds to one decimal. Each coin adds a tiny rounding error. After 500 coins, you might be off by one. But the scale reads twice as fast. FP16 does math at 2x the speed of FP32, but accumulates small rounding errors over hundreds of additions.
|
| 68 |
+
|
| 69 |
+
**Technical:** The dequant matvec inner loop changes from float32 to float16 accumulation. Apple's A-series GPU has dedicated fp16 ALUs at 2x throughput. The FMA becomes half-precision: `fma(half(nibble), half(scale*x), half(bias*x))`. Final output is promoted to float32 via `simd_sum`. Risk: fp16 has ~3 decimal digits; sums of 512+ elements may lose precision.
|
| 70 |
+
|
| 71 |
+
**Default:** OFF (experimental)
|
| 72 |
+
|
| 73 |
+
---
|
| 74 |
+
|
| 75 |
+
## FP8 KV Cache
|
| 76 |
+
|
| 77 |
+
**Analogy:** The KV cache is like a notebook where the model writes down what it's seen. FP32 uses a full page per note. FP8 uses a quarter page -- same content, just more compressed handwriting. You fit 4x more notes in the same notebook, so the model can remember 4x more conversation.
|
| 78 |
+
|
| 79 |
+
**Technical:** Stores attention Key and Value vectors in FP8 E4M3 format (1 byte vs 4 bytes per element) with per-position dynamic scaling. Encoding: absmax/240 scale factor, each float clipped and quantized to 8-bit (1 sign, 4 exponent, 3 mantissa). Decoding is inline during attention compute via Metal function constants. 4x memory reduction enables 4x longer context at the same memory budget.
|
| 80 |
+
|
| 81 |
+
**Default:** OFF
|
| 82 |
+
|
| 83 |
+
See [context-optimization.md](context-optimization.md) for detailed memory savings tables.
|
| 84 |
+
|
| 85 |
+
---
|
| 86 |
+
|
| 87 |
+
## Max Context Length
|
| 88 |
+
|
| 89 |
+
**Analogy:** Context length is how far back the model can "see" in the conversation. Like a person's short-term memory -- 4K tokens is the last few minutes, 32K is the last hour. More context = better understanding of the conversation, but uses more memory.
|
| 90 |
+
|
| 91 |
+
**Technical:** Sets the maximum sequence length for KV cache allocation. Memory cost: `num_full_attn_layers x 2 (K+V) x kv_heads x head_dim x bytes_per_elem x positions`. For the 35B with 10 full-attn layers: 40KB/pos (FP32) or 10KB/pos (FP8). Auto mode uses `os_proc_available_memory()` to pick the largest safe value.
|
| 92 |
+
|
| 93 |
+
**Options:** Auto, 4K, 8K, 16K, 32K
|
| 94 |
+
|
| 95 |
+
---
|
| 96 |
+
|
| 97 |
+
## Sliding Window
|
| 98 |
+
|
| 99 |
+
**Analogy:** Instead of remembering everything forever (which fills up memory), the full-attention layers only look at the last N tokens -- like a window sliding along the conversation. But the 30 linear attention layers still remember everything through their state matrices. It's like having both short-term and long-term memory working together.
|
| 100 |
+
|
| 101 |
+
**Technical:** Implements a circular KV buffer for full attention layers. Write: `cache_pos = kv->len % window_size`. Read: attend only to the most recent `window_size` positions. The 30 GatedDeltaNet layers maintain full context via their 128x128 state matrices (O(1) memory). Only the 10 full attention layers are windowed. With window 4096 + FP8: fixed 40MB KV regardless of conversation length.
|
| 102 |
+
|
| 103 |
+
**Options:** Off, 2048, 4096, 8192
|
| 104 |
+
|
| 105 |
+
See [context-optimization.md](context-optimization.md) for the full context optimization story.
|
| 106 |
+
|
| 107 |
+
---
|
| 108 |
+
|
| 109 |
+
## Thinking Mode
|
| 110 |
+
|
| 111 |
+
**Analogy:** Like a student who shows their work before giving the final answer. The model reasons step-by-step inside `<think>` tags before responding. This usually produces better answers, but takes more tokens (and time). At low K values, the model may get stuck thinking forever -- disable it for speed.
|
| 112 |
+
|
| 113 |
+
**Technical:** The chat template includes a `<think>` tag after the assistant turn header. The model generates reasoning tokens inside the think block, then emits `</think>` before the actual response. Think budget caps the maximum thinking tokens and force-emits `</think>`. Set to -1 to disable thinking entirely (removes `<think>` from the template).
|
| 114 |
+
|
| 115 |
+
**Default:** ON (with configurable budget, default 2048 tokens)
|
| 116 |
+
|
| 117 |
+
---
|
| 118 |
+
|
| 119 |
+
## H2O Budget (Coming Soon)
|
| 120 |
+
|
| 121 |
+
**Analogy:** Instead of sliding window's simple "forget old stuff" approach, H2O watches which parts of the conversation the model keeps looking back at (the "heavy hitters") and keeps those. It's like a librarian who notices which reference books get used most and keeps those on the desk, while shelving the rarely-used ones.
|
| 122 |
+
|
| 123 |
+
**Technical:** Heavy Hitter Oracle eviction policy. Tracks cumulative post-softmax attention scores per KV position. Budget = total positions to keep. Protected regions: sink tokens (first 4) + recent tokens (25% of budget). Remaining budget goes to positions with highest cumulative attention scores. After eviction, both CPU and GPU caches are compacted to contiguous positions. Replaces sliding window when both are configured (H2O is strictly better).
|
| 124 |
+
|
| 125 |
+
**Status:** Implementation complete in the engine (`kv_cache_evict_h2o`, `kv_cache_h2o_accumulate_scores`). UI toggle pending.
|
docs/io-and-gpu-exploration.md
ADDED
|
@@ -0,0 +1,215 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# I/O and GPU Exploration: What We Learned Running a 397B Model from SSD
|
| 2 |
+
|
| 3 |
+
## The Problem
|
| 4 |
+
|
| 5 |
+
We're streaming a 397 billion parameter Mixture-of-Experts model from NVMe SSD on a MacBook Pro with 48GB RAM. The model's expert weights total 120GB at 2-bit quantization (209GB at 4-bit). Only 6GB fits in memory. Every generated token requires reading ~600MB of expert data from disk — 4 experts × 3.9MB × 60 layers.
|
| 6 |
+
|
| 7 |
+
The question that drove months of optimization: **where does the time actually go, and what can we do about it?**
|
| 8 |
+
|
| 9 |
+
## Part 1: The GPU Story
|
| 10 |
+
|
| 11 |
+
### What the profiler showed us
|
| 12 |
+
|
| 13 |
+
We captured a Metal GPU trace of the expert forward pass (the most compute-intensive per-token operation). The results were surprising:
|
| 14 |
+
|
| 15 |
+
- **Total GPU compute time for 20 expert matvecs: 747µs** (37µs each)
|
| 16 |
+
- **Wall clock time: 31.5ms** (1.58ms each)
|
| 17 |
+
- **GPU utilization: 2.4%**
|
| 18 |
+
|
| 19 |
+
The GPU finishes its actual math in microseconds, then sits idle waiting for the next batch of data. The "compute" bottleneck is really a **data delivery** bottleneck.
|
| 20 |
+
|
| 21 |
+
### Where GPU cycles go
|
| 22 |
+
|
| 23 |
+
The instruction cost breakdown from Metal's performance counters:
|
| 24 |
+
|
| 25 |
+
| Category | % of GPU Time | What it is |
|
| 26 |
+
|----------|:---:|---|
|
| 27 |
+
| Math (FMA/MUL/ADD) | 63.9% | The actual dequant + multiply-accumulate |
|
| 28 |
+
| Conversion | 25.0% | `bf16_to_f32()` calls for scale/bias lookup |
|
| 29 |
+
| Data Movement | 3.3% | Moving data between register files |
|
| 30 |
+
| Bit Manipulation | 19.0% | Extracting 2-bit values from uint32 |
|
| 31 |
+
|
| 32 |
+
**25% of GPU time is type conversion.** Every scale and bias value is stored as bfloat16 and converted to float32 on the fly. Storing them as float32 would double the scale/bias storage (negligible overall) but eliminate a quarter of the GPU work.
|
| 33 |
+
|
| 34 |
+
**19% is bit manipulation.** Extracting 16 × 2-bit values from each uint32 requires shift-and-mask operations. A lookup table or wider SIMD approach could reduce this.
|
| 35 |
+
|
| 36 |
+
But here's the kicker: even if we eliminated ALL of this overhead, saving 44% of GPU time would save 44% of 37µs = 16µs per matvec. At 720 matvecs per token (4 experts × 3 projections × 60 layers), that's 11.5ms per token. Not nothing — but the I/O bottleneck at 90ms/token makes it a secondary concern.
|
| 37 |
+
|
| 38 |
+
### Cache behavior
|
| 39 |
+
|
| 40 |
+
- **L1 Cache Read Hit Rate: 93.4%** — our threadgroup shared memory (`x_shared[4096]`) works perfectly for caching the input vector
|
| 41 |
+
- **L1 Cache Write Hit Rate: 100%** — all output writes hit L1
|
| 42 |
+
- **Last Level Cache Bandwidth: ~418 GB/s** — nearly saturating unified memory bandwidth
|
| 43 |
+
|
| 44 |
+
The GPU cache hierarchy is working well. The 6.6% L1 read miss rate corresponds to expert weight data that doesn't fit in L1 (each expert is 3.9MB, L1 is ~192KB per core). These misses go to the shared L2 / memory fabric, which runs at near-theoretical bandwidth.
|
| 45 |
+
|
| 46 |
+
### The GPU cluster affinity experiment
|
| 47 |
+
|
| 48 |
+
Apple's M3 Max has 40 GPU cores organized into clusters, each with its own L2 cache (~4MB). Our 2-bit experts are 3.9MB — almost exactly one cluster's L2 capacity.
|
| 49 |
+
|
| 50 |
+
**Hypothesis:** If we encode all 4 operations for one expert (gate → up → SwiGLU → down) into a single Metal command encoder, the GPU scheduler would keep that work on one cluster, and the expert's weight data would stay hot in that cluster's L2.
|
| 51 |
+
|
| 52 |
+
**Result:** 2% slower. The fused single-encoder approach reduced parallelism — Metal's scheduler couldn't overlap work across experts anymore. The existing 2-encoder-per-expert approach (gate+up together, SwiGLU+down together) lets the GPU interleave expert computations across clusters, which provides better throughput than L2 locality.
|
| 53 |
+
|
| 54 |
+
**Lesson:** GPU schedulers are smarter than manual NUMA pinning. Don't fight the hardware scheduler unless you have profiling data showing it's making bad decisions.
|
| 55 |
+
|
| 56 |
+
### What doesn't matter on the GPU
|
| 57 |
+
|
| 58 |
+
- **Superpages (2MB pages):** Apple Silicon ARM64 uses fixed 16KB pages. `vm_allocate` with `VM_FLAGS_SUPERPAGE_SIZE_2MB` returns `KERN_INVALID_ARGUMENT`. Not available.
|
| 59 |
+
- **Command buffer type:** `commandBufferWithUnretainedReferences` (skip ARC retain/release) vs `commandBuffer` — zero measurable difference.
|
| 60 |
+
- **Encoder count:** Batching all experts' gate+up into one encoder vs separate encoders per expert — zero difference. Metal handles both patterns efficiently.
|
| 61 |
+
|
| 62 |
+
## Part 2: The I/O Story
|
| 63 |
+
|
| 64 |
+
### The landscape
|
| 65 |
+
|
| 66 |
+
Our I/O benchmark measured raw SSD performance for the expert read pattern:
|
| 67 |
+
|
| 68 |
+
| Access Pattern | Throughput | Latency (4 experts) |
|
| 69 |
+
|---|:---:|:---:|
|
| 70 |
+
| Sequential, warm page cache | 32.1 GB/s | 0.49 ms |
|
| 71 |
+
| Parallel 4T, warm cache | 29.2 GB/s | 0.97 ms |
|
| 72 |
+
| Parallel 4T, cold (F_NOCACHE) | 5.5 GB/s | 2.84 ms |
|
| 73 |
+
| Sequential, cold | 4.5 GB/s | 3.46 ms |
|
| 74 |
+
| mmap + memcpy, cold | 0.12 GB/s | varies |
|
| 75 |
+
|
| 76 |
+
The gap between warm (32 GB/s) and cold (5.5 GB/s) is the entire optimization story. Everything we tried was about moving more data from cold to warm.
|
| 77 |
+
|
| 78 |
+
### The mmap disaster
|
| 79 |
+
|
| 80 |
+
**What:** Replace `pread()` with `mmap()` + `memcpy()` for zero-syscall access to cached data.
|
| 81 |
+
|
| 82 |
+
**Result:** 0.56 tok/s — **5x slower** than pread.
|
| 83 |
+
|
| 84 |
+
**Why:** Each 3.9MB expert spans 240 × 16KB pages. For uncached data, mmap triggers 240 individual page faults, each requiring a separate kernel trap → I/O request → page table update. A single `pread()` call issues one large NVMe command for the entire 3.9MB range.
|
| 85 |
+
|
| 86 |
+
**Lesson:** `mmap()` is designed for random access to already-cached data. For bulk reads of potentially uncached data, `pread()` is dramatically better because it lets the kernel optimize the I/O pattern.
|
| 87 |
+
|
| 88 |
+
### The custom cache trap
|
| 89 |
+
|
| 90 |
+
We built increasingly sophisticated expert caching systems:
|
| 91 |
+
|
| 92 |
+
| Cache Type | Entries | Memory | Hit Rate | tok/s | Verdict |
|
| 93 |
+
|---|:---:|:---:|:---:|:---:|---|
|
| 94 |
+
| None (pread only) | 0 | 0 | 0% | 2.86 | Baseline |
|
| 95 |
+
| Metal LRU (500) | 500 | 3.5 GB | 35% | 3.14 | Small win |
|
| 96 |
+
| Metal LRU (1000) | 1000 | 7.1 GB | 44% | 2.24 | **Worse** |
|
| 97 |
+
| Metal LRU (2500) | 2500 | 9.8 GB | 55% | 2.24 | **Worse** |
|
| 98 |
+
| Metal LRU (3000) | 3000 | 21 GB | 55% | 1.99 | **Much worse** |
|
| 99 |
+
| Malloc zero-copy (2581) | 2581 | 18 GB | 52% | 2.10 | **Worse** |
|
| 100 |
+
| **No cache, trust OS** | **0** | **0** | **OS-managed** | **5.74** | **Best** |
|
| 101 |
+
|
| 102 |
+
**The breakthrough:** Deleting the entire custom cache system and letting macOS manage the page cache yielded a **38% speedup** over our best custom implementation.
|
| 103 |
+
|
| 104 |
+
**Why custom caches hurt:**
|
| 105 |
+
1. **Metal buffer caches wire memory.** Every Metal buffer allocation is pinned in physical RAM (wired pages). Our 9.8GB cache wired 9.8GB, leaving only ~25GB for the OS page cache instead of ~35GB.
|
| 106 |
+
2. **The OS page cache is smarter.** macOS uses CLOCK-Pro (an adaptive replacement algorithm that balances recency and frequency). Our LRU cache was strictly recency-based.
|
| 107 |
+
3. **Zero lookup overhead.** The OS page cache operates at the virtual memory level — a cache "hit" is just a normal memory access through the MMU. Our cache required hash table lookups, pointer chasing, and LRU bookkeeping.
|
| 108 |
+
4. **Memory pressure compounds.** `vm_stat` monitoring showed that with the Metal cache active, the compressor was doing 60,000-130,000 decompressions per second. Without it: near zero. The wired cache pages forced the OS to compress other data, and decompressing it on access added latency everywhere.
|
| 109 |
+
|
| 110 |
+
**The database analogy:** PostgreSQL recommends keeping `shared_buffers` at 25% of RAM and letting the OS cache handle the rest. We were doing the equivalent of setting shared_buffers to 60% of RAM — squeezing out the OS cache that handles the long tail of access patterns better than any application-level cache.
|
| 111 |
+
|
| 112 |
+
### The kernel hint experiments
|
| 113 |
+
|
| 114 |
+
We tried every macOS I/O hint available:
|
| 115 |
+
|
| 116 |
+
| Hint | Purpose | Result | Why |
|
| 117 |
+
|---|---|:---:|---|
|
| 118 |
+
| `F_NOCACHE` | Bypass page cache | +3% (2-bit) | Avoids thrashing when working set >> cache. But prevents warm hits. |
|
| 119 |
+
| `F_RDAHEAD` | Enable readahead | 0% | Kernel already does readahead for pread. |
|
| 120 |
+
| `F_RDADVISE` (immediate) | Pre-hint reads | -8% | Creates NVMe command contention — double-issues reads. |
|
| 121 |
+
| `F_RDADVISE` (with lead time) | Pre-hint from previous token | -4% | 65-80% of predictions wrong (different routing). Wrong advises waste bandwidth. |
|
| 122 |
+
| `MADV_RANDOM` | Disable readahead | **Harmful** | Fragments 3.9MB reads into 5.7 × 512KB disk ops. |
|
| 123 |
+
| `MADV_SEQUENTIAL` | Large readahead | 0% | Fragmentation is physical page layout, not readahead policy. |
|
| 124 |
+
| `MADV_WILLNEED` | Pre-populate cache | 0% on steady state | Only helps first access, not sustained generation. |
|
| 125 |
+
| No hint (default) | Let kernel decide | **Best** | Kernel's default behavior is already well-tuned for Apple hardware. |
|
| 126 |
+
|
| 127 |
+
**The pattern:** Every hint we tried either made no difference or made things worse. The macOS kernel is already optimized for Apple's NVMe controller. Application-level hints add overhead (each `fcntl` / `madvise` is a syscall) without providing information the kernel doesn't already have.
|
| 128 |
+
|
| 129 |
+
### The fragmentation discovery
|
| 130 |
+
|
| 131 |
+
`fs_usage` profiling revealed the kernel's internal behavior:
|
| 132 |
+
|
| 133 |
+
```
|
| 134 |
+
pread calls: 45,414
|
| 135 |
+
RdData ops: 260,845
|
| 136 |
+
Reads per pread: 5.7x
|
| 137 |
+
```
|
| 138 |
+
|
| 139 |
+
Each 3.9MB pread is broken into ~5.7 separate NVMe commands, mostly 512KB (0x80000) and smaller. The reason: the page cache stores data in scattered 16KB virtual pages that map to non-contiguous physical pages. The kernel can't coalesce them into a single DMA transfer.
|
| 140 |
+
|
| 141 |
+
**Block size distribution from fs_usage:**
|
| 142 |
+
```
|
| 143 |
+
76,549 × 512KB (0x80000)
|
| 144 |
+
38,441 × 8KB (0x2000)
|
| 145 |
+
20,064 × 16KB (0x4000)
|
| 146 |
+
17,647 × 12KB (0x3000)
|
| 147 |
+
14,651 × 256KB (0x40000)
|
| 148 |
+
```
|
| 149 |
+
|
| 150 |
+
This fragmentation adds ~46µs of kernel overhead per pread (240 pread calls/token × 46µs = 11ms/token). It's inherent to the virtual memory system and can't be fixed from userspace.
|
| 151 |
+
|
| 152 |
+
### Buffer alignment matters
|
| 153 |
+
|
| 154 |
+
Our isolated benchmark showed a dramatic difference based on destination buffer alignment:
|
| 155 |
+
|
| 156 |
+
| Buffer Alignment | Avg Latency | Throughput |
|
| 157 |
+
|---|:---:|:---:|
|
| 158 |
+
| 2MB-aligned (`posix_memalign`) | 234 µs | 16.8 GB/s |
|
| 159 |
+
| 16KB-aligned (default Metal) | 836 µs | 4.7 GB/s |
|
| 160 |
+
|
| 161 |
+
**3.6x faster with 2MB alignment** for page-cache-resident data. The DMA controller can do larger, more efficient burst transfers when the destination is aligned to large boundaries.
|
| 162 |
+
|
| 163 |
+
In the full pipeline, the improvement was more modest (5%) because most reads hit SSD (cold data), where the DMA controller's performance is dominated by NAND flash latency rather than buffer alignment. But it's a free optimization — `posix_memalign` + `newBufferWithBytesNoCopy` costs nothing at runtime.
|
| 164 |
+
|
| 165 |
+
### The tiered I/O experiment
|
| 166 |
+
|
| 167 |
+
**Hypothesis:** Use two file descriptors per layer file — one with `F_NOCACHE` for first-time reads (avoid polluting page cache with one-off data), one without for repeat reads (benefit from page cache). Track "seen" experts with a 3.8KB bitset.
|
| 168 |
+
|
| 169 |
+
**Result:** Marginally better tok/s, but `vm_stat` showed identical memory pressure. The memory pressure was from Metal buffers and model weights, not from page cache behavior. The tiered approach added complexity without meaningful benefit.
|
| 170 |
+
|
| 171 |
+
### What actually worked for I/O
|
| 172 |
+
|
| 173 |
+
1. **2-bit expert quantization** — 44% smaller files. Reduced expert_io from 2.6ms to 1.5ms per layer. The single biggest improvement.
|
| 174 |
+
2. **Trust the OS page cache** — Delete custom caches. Let macOS manage memory. 38% speedup.
|
| 175 |
+
3. **2MB-aligned DMA buffers** — 5% improvement on expert_io. Free optimization.
|
| 176 |
+
4. **Parallel pread (4 threads)** — 9.2x speedup over sequential (superlinear due to NVMe command queuing).
|
| 177 |
+
5. **No kernel hints** — Default behavior is already optimal. Every hint we tried was neutral or harmful.
|
| 178 |
+
|
| 179 |
+
## Part 3: The Bigger Picture
|
| 180 |
+
|
| 181 |
+
### The SSD bandwidth wall
|
| 182 |
+
|
| 183 |
+
At 2-bit precision, the theoretical I/O floor for our workload is:
|
| 184 |
+
|
| 185 |
+
```
|
| 186 |
+
60 layers × ~2.6 cache misses × 3.9MB = 608MB per token
|
| 187 |
+
608MB ÷ 5.5 GB/s (random read throughput) = 110ms
|
| 188 |
+
110ms → 9.1 tok/s theoretical maximum (I/O limited)
|
| 189 |
+
```
|
| 190 |
+
|
| 191 |
+
Our measured performance of 5.5 tok/s (182ms/token) is split roughly 50/50 between I/O (90ms) and compute (90ms). We're at 82% of the I/O-limited theoretical maximum.
|
| 192 |
+
|
| 193 |
+
The remaining 18% gap is the page cache fragmentation overhead (5.7 ops/pread) and the kernel's per-read overhead (46µs/pread). These are architectural limitations of macOS's virtual memory system.
|
| 194 |
+
|
| 195 |
+
### Systems thinking beats micro-optimization
|
| 196 |
+
|
| 197 |
+
The single most impactful change in this entire project was **deleting code**: removing the 9.8GB Metal buffer cache. It wasn't that the cache was poorly implemented — it was that the cache's existence created system-level effects (memory pressure, compressor thrashing, reduced page cache) that outweighed its direct benefits.
|
| 198 |
+
|
| 199 |
+
This is a classic systems engineering lesson: optimizing one component in isolation can degrade the whole system. The GPU profiling showed us the compute isn't the bottleneck. The I/O profiling showed us the kernel is already doing a good job. The `vm_stat` monitoring showed us our "optimization" was causing the real problem.
|
| 200 |
+
|
| 201 |
+
### What the database world already knew
|
| 202 |
+
|
| 203 |
+
Dan Woods brought the key insight: **treat the model weights like a database.** Databases have solved the problem of accessing datasets larger than memory for decades:
|
| 204 |
+
|
| 205 |
+
- **Don't build your own buffer pool.** PostgreSQL learned this — `shared_buffers` should be 25% of RAM, not 100%. The OS buffer cache handles the long tail better.
|
| 206 |
+
- **Respect the hardware cache hierarchy.** Don't bypass caches (F_NOCACHE) unless you have measured evidence of thrashing. The caches exist for a reason.
|
| 207 |
+
- **Profile before optimizing.** `fs_usage` and Metal GPU traces told us exactly where time goes. Without them, we would have optimized the wrong thing.
|
| 208 |
+
- **Alignment matters for DMA.** Database systems align I/O buffers to page boundaries. We found 2MB alignment gives 3.6x better DMA throughput.
|
| 209 |
+
|
| 210 |
+
### Remaining frontiers
|
| 211 |
+
|
| 212 |
+
1. **Batch prefill** — Process multiple prompt tokens simultaneously. The sequential GatedDeltaNet recurrence limits parallelism, but projection matmuls and expert I/O can be batched.
|
| 213 |
+
2. **C tokenizer** — Done. Eliminated the 3.5s Python overhead, bringing setup from 4s to 180ms.
|
| 214 |
+
3. **The page cache fragmentation** — 5.7 disk ops per pread is a kernel limitation. The only userspace mitigation would be to use `mincore()` to detect cached pages and `memcpy()` from mmap for hits, falling back to pread for misses.
|
| 215 |
+
4. **Expert file layout optimization** — Co-locating frequently co-accessed experts could reduce the number of distinct NVMe commands per token. This requires offline analysis of routing patterns.
|
docs/ios-port.md
ADDED
|
@@ -0,0 +1,244 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# iOS Port Overview
|
| 2 |
+
|
| 3 |
+
This document summarizes the Flash-MoE iOS port. For full details, see:
|
| 4 |
+
- [FlashMoE-iOS/IOS_PORT.md](../FlashMoE-iOS/IOS_PORT.md) -- complete porting story, problems solved, architecture
|
| 5 |
+
- [FlashMoE-iOS/397B_ANALYSIS.md](../FlashMoE-iOS/397B_ANALYSIS.md) -- 397B on iPhone: memory budget, Metal limits, K-reduction quality
|
| 6 |
+
|
| 7 |
+
## What We Built
|
| 8 |
+
|
| 9 |
+
A native SwiftUI iOS app that runs Qwen3.5 MoE models on iPhone, sharing 100% of the C/Metal inference engine with the macOS CLI via unity build (`#include "infer.m"`).
|
| 10 |
+
|
| 11 |
+
### Results
|
| 12 |
+
|
| 13 |
+
| Device | Model | K | tok/s | Notes |
|
| 14 |
+
|--------|-------|---|-------|-------|
|
| 15 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-35B-A3B | 8 | **5.5** | Full quality, full GPU path |
|
| 16 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-35B-A3B (tiered) | 8 | **5.5+** | 13.4GB download, same quality |
|
| 17 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-35B-A3B | 4 | **~11*** | *Projected with K=4 + all optimizations (CMD merge, fused expert, prefetch) |
|
| 18 |
+
| iPhone 17 (12GB, A19) | Qwen3.5-397B-A17B | 4 | ~0.003 | CPU fallback only (Metal 4GB buffer limit) |
|
| 19 |
+
|
| 20 |
+
iPhone achieves 57% of laptop speed on the 35B model with 17% of the memory. With K=4 and all Expert Settings optimizations enabled, projections reach ~11 tok/s on the 35B.
|
| 21 |
+
|
| 22 |
+
## iOS App Features
|
| 23 |
+
|
| 24 |
+
### Chat Interface
|
| 25 |
+
- Streaming token display with typing animation
|
| 26 |
+
- Message bubbles with text selection (long press to copy)
|
| 27 |
+
- Collapsible `<think>` blocks (DisclosureGroup)
|
| 28 |
+
- Special token stripping (`<|endoftext|>`, `<|im_end|>`, `<|im_start|>`)
|
| 29 |
+
- Tap outside keyboard to dismiss
|
| 30 |
+
- Auto-scroll to latest message via `ScrollViewReader` + `scrollTo(:anchor:.bottom)`
|
| 31 |
+
- New chat / reset conversation
|
| 32 |
+
- KV cache reuse across conversation turns (continuation mode)
|
| 33 |
+
|
| 34 |
+
### Model Management
|
| 35 |
+
- On-device model scanning (Documents directory)
|
| 36 |
+
- HuggingFace download catalog with per-model K recommendations
|
| 37 |
+
- Background URLSession downloads with progress tracking
|
| 38 |
+
- Swipe-to-delete for downloaded models
|
| 39 |
+
- Import from Files app (UIDocumentPickerViewController) with bookmark or move-to-Documents
|
| 40 |
+
- Export/Move model to Files app for cross-app access
|
| 41 |
+
- Model info sheet (layers, experts, hidden dim, vocab, file sizes)
|
| 42 |
+
|
| 43 |
+
### Expert Settings
|
| 44 |
+
All settings include info modals with plain-language analogies and technical explanations. Compact UI layout with info icon to the left of each label.
|
| 45 |
+
|
| 46 |
+
- **Active Experts (K)** — K value picker (2-10) for K-reduction (fewer experts = less I/O, lower quality)
|
| 47 |
+
- **I/O Fanout** — Chunks picker (off/2/4/8) for splitting expert reads into parallel chunks
|
| 48 |
+
- **CMD1+CMD2 Merge** — Combine GPU command buffers for linear attention layers (saves one sync per layer)
|
| 49 |
+
- **Fused Attention** — Single-kernel FlashAttention-style online softmax (replaces 3-dispatch pipeline)
|
| 50 |
+
- **Fused Expert Kernel** — Combined gate+up+SwiGLU in one Metal dispatch per expert
|
| 51 |
+
- **Expert Prefetch** — Overlap next-layer expert I/O with current-layer GPU compute
|
| 52 |
+
- **FP16 Accumulation** — Experimental half-precision accumulation in dequant kernels (default OFF)
|
| 53 |
+
- **FP8 KV Cache** — 4x KV memory reduction via FP8 E4M3 quantization
|
| 54 |
+
- **Max Context Length** — Selector from 4K to 32K positions (Auto mode uses `os_proc_available_memory()`)
|
| 55 |
+
- **Sliding Window** — Circular KV buffer for full attention layers (0/2048/4096/8192)
|
| 56 |
+
- **Thinking Mode** — Enable/disable `<think>` chain-of-thought with configurable budget
|
| 57 |
+
- **H2O Budget** — Heavy Hitter Oracle KV eviction (coming soon)
|
| 58 |
+
- Max generation tokens bumped from 500 to 2048
|
| 59 |
+
|
| 60 |
+
### Profiler
|
| 61 |
+
- Resource monitoring overlay
|
| 62 |
+
- Thermal state indicator (Cool/Warm/Hot/Critical)
|
| 63 |
+
- Temperature display in Celsius
|
| 64 |
+
- TTFT display in minutes when >500s
|
| 65 |
+
- tok/s and tokens generated counters
|
| 66 |
+
|
| 67 |
+
### Quantization Support
|
| 68 |
+
- 4-bit experts (full quality, production)
|
| 69 |
+
- 2-bit experts (faster, breaks JSON/tool calling)
|
| 70 |
+
- Tiered quantization (4-bit hot / 2-bit cold experts, auto-detected)
|
| 71 |
+
|
| 72 |
+
### Context Management
|
| 73 |
+
- **FP8 KV Cache** — FP8 E4M3 quantization: float32 (4 bytes) to uint8 (1 byte) per element. Per-position dynamic scales in separate Metal buffers. 4x memory reduction (~60KB to ~15KB per position for 397B). GPU inline dequant in fused attention kernel. FP8/sliding window flags are set BEFORE `metal_setup()` to ensure correct buffer allocation.
|
| 74 |
+
- **Sliding Window Attention** — Circular KV buffer for full attention layers. Only the 10 full attention layers are windowed; the 30 GatedDeltaNet layers maintain full context via 128x128 state matrices. With window 4096 + FP8: fixed 40MB KV regardless of conversation length.
|
| 75 |
+
- **Max Context Length** — Configurable from 4K to 32K. Auto mode uses `os_proc_available_memory()` to pick the largest safe value. Memory cost: num_full_attn_layers x 2 x kv_heads x head_dim x bytes_per_elem x positions.
|
| 76 |
+
- **H2O KV Cache Eviction** (in progress) — Heavy Hitter Oracle: tracks cumulative attention scores, protects sink tokens + recent tokens, evicts low-scoring positions. Replaces sliding window when both are configured. See [context-optimization.md](context-optimization.md).
|
| 77 |
+
|
| 78 |
+
### Custom URL Download
|
| 79 |
+
- Paste any HuggingFace model URL (e.g. `mlx-community/Qwen3.5-35B-A3B-4bit`) in the download section
|
| 80 |
+
- URL is validated and config.json is fetched to verify Qwen3.5 MoE compatibility
|
| 81 |
+
- Custom models appear in the download list alongside catalog entries
|
| 82 |
+
- Downloaded models are hidden from the catalog (no duplicate entries)
|
| 83 |
+
- Trash icon removed from catalog download rows
|
| 84 |
+
|
| 85 |
+
### Universal App Support
|
| 86 |
+
- Same SwiftUI shell compiles for both iPhone and Mac destinations
|
| 87 |
+
- `#if os(iOS)` conditional compilation for platform-specific UI (toolbar, keyboard dismiss, document picker)
|
| 88 |
+
- C inference engine, Metal shaders, and Swift bridge are fully cross-platform
|
| 89 |
+
- No code fork — one codebase serves both platforms
|
| 90 |
+
- macOS sandbox entitlements: app-sandbox, extended-virtual-addressing, increased-memory-limit, user-selected file read-write, network client
|
| 91 |
+
|
| 92 |
+
### Model Management
|
| 93 |
+
- On-device model scanning (Documents directory)
|
| 94 |
+
- HuggingFace download catalog with per-model K recommendations
|
| 95 |
+
- Background URLSession downloads with progress tracking
|
| 96 |
+
- Swipe-to-delete for downloaded models
|
| 97 |
+
- Import from Files app (UIDocumentPickerViewController) with bookmark or move-to-Documents
|
| 98 |
+
- Export/Move model to Files app for cross-app access
|
| 99 |
+
- Model info sheet (layers, experts, hidden dim, vocab, file sizes)
|
| 100 |
+
- Custom HuggingFace URL download (see above)
|
| 101 |
+
|
| 102 |
+
## Architecture
|
| 103 |
+
|
| 104 |
+
```
|
| 105 |
+
SwiftUI (UI + @Observable state)
|
| 106 |
+
-> Swift async bridge (AsyncStream<GenerationToken>)
|
| 107 |
+
-> Objective-C wrapper (FlashMoEEngine.h C API)
|
| 108 |
+
-> C inference engine (7,500+ lines, unity build)
|
| 109 |
+
-> Metal GPU shaders (1,300 lines)
|
| 110 |
+
```
|
| 111 |
+
|
| 112 |
+
### Engine C API (FlashMoEEngine.h)
|
| 113 |
+
|
| 114 |
+
```
|
| 115 |
+
flashmoe_create() -- allocate context
|
| 116 |
+
flashmoe_load(ctx, config) -- load model, allocate Metal resources
|
| 117 |
+
flashmoe_generate(ctx, prompt, max_tokens, callback, user_data)
|
| 118 |
+
flashmoe_generate_continuation() -- reuse KV cache for multi-turn
|
| 119 |
+
flashmoe_cancel() -- thread-safe cancellation
|
| 120 |
+
flashmoe_reset() -- clear KV cache and position
|
| 121 |
+
flashmoe_unload() -- release model resources
|
| 122 |
+
flashmoe_destroy() -- free context
|
| 123 |
+
flashmoe_get_stats() -- model info + generation stats
|
| 124 |
+
flashmoe_validate_model() -- check model directory validity
|
| 125 |
+
flashmoe_turn_count() -- conversation turn count
|
| 126 |
+
flashmoe_last_error() -- human-readable error string
|
| 127 |
+
```
|
| 128 |
+
|
| 129 |
+
### Key Files
|
| 130 |
+
|
| 131 |
+
```
|
| 132 |
+
FlashMoE-iOS/
|
| 133 |
+
FlashMoEEngine/
|
| 134 |
+
FlashMoEEngine.h -- C API (create/load/generate/cancel/reset/destroy)
|
| 135 |
+
FlashMoEEngine.m -- Unity build wrapping infer.m (#define CHAT_MODE 1)
|
| 136 |
+
Bridge/
|
| 137 |
+
FlashMoEBridge.swift -- @Observable async Swift wrapper
|
| 138 |
+
Views/
|
| 139 |
+
ChatView.swift -- Streaming chat UI with thinking disclosure
|
| 140 |
+
ModelListView.swift -- Model discovery + download catalog
|
| 141 |
+
ModelDownloadRow.swift -- Download progress with pause/resume
|
| 142 |
+
ProfilerView.swift -- Resource monitoring overlay
|
| 143 |
+
Services/
|
| 144 |
+
DownloadManager.swift -- Background URLSession model downloads
|
| 145 |
+
Models/
|
| 146 |
+
ModelCatalog.swift -- HuggingFace model registry with K recommendations
|
| 147 |
+
App/
|
| 148 |
+
FlashMoEApp.swift -- SwiftUI app entry point
|
| 149 |
+
IOS_PORT.md -- Full porting documentation
|
| 150 |
+
397B_ANALYSIS.md -- 397B memory/performance analysis
|
| 151 |
+
project.yml -- XcodeGen config (iOS 18+, iPhone only)
|
| 152 |
+
copy_model_to_iphone.sh -- Push models to device over USB
|
| 153 |
+
```
|
| 154 |
+
|
| 155 |
+
## iOS-Specific Constraints and Solutions
|
| 156 |
+
|
| 157 |
+
### Metal 4GB Per-Buffer Limit
|
| 158 |
+
|
| 159 |
+
iOS Metal buffers cannot exceed 4096 MB regardless of entitlements. The 35B model weights (~2.5GB) fit; the 397B weights (~5.5GB) do not.
|
| 160 |
+
|
| 161 |
+
**Attempted workarounds (all failed):**
|
| 162 |
+
- Single 5.5GB Metal buffer -- Metal assertion crash
|
| 163 |
+
- Two overlapping ~3GB Metal buffers -- OOM kill (8GB shared memory on 12GB device)
|
| 164 |
+
- 50MB staging buffer with memcpy per dispatch -- data corruption from in-flight command buffer aliasing
|
| 165 |
+
- CPU fallback -- works but 6 min/token
|
| 166 |
+
|
| 167 |
+
**Solution:** Split `model_weights.bin` into two <4GB files at the Python packing stage (pending implementation).
|
| 168 |
+
|
| 169 |
+
### Memory Management
|
| 170 |
+
- Adaptive context length via `os_proc_available_memory()` -- reduces 262144 to 8192 based on available memory
|
| 171 |
+
- Wired memory budget: `recommendedMaxWorkingSetSize` constrains Metal buffer totals to stay within device budget
|
| 172 |
+
- KV cache sizing: `MAX_SEQ_LEN` (1M) replaced with runtime `g_kv_seq_len` (4096) per cache
|
| 173 |
+
- FP8 KV cache: 4x memory reduction when opted in (`bytes_per_elem = g_use_fp8_kv ? 1 : sizeof(float)`)
|
| 174 |
+
- Expert mmap disabled on iOS -- jetsam kills from 112GB mapped address space
|
| 175 |
+
- Debug vs Release: Metal debug wrappers add ~2GB overhead, must build Release for on-device testing
|
| 176 |
+
- `isExcludedFromBackup` on all model files to prevent iOS purging 200GB+ of data
|
| 177 |
+
|
| 178 |
+
### OOM Prevention (8 Protections)
|
| 179 |
+
|
| 180 |
+
1. **Memory pressure dispatch source** — `DISPATCH_SOURCE_TYPE_MEMORYPRESSURE` handler cancels generation on `DISPATCH_MEMORYPRESSURE_CRITICAL`. This is actionable (sets `atomic_store(&ctx->cancelled, 1)`), not just logging.
|
| 181 |
+
2. **`didReceiveMemoryWarning` observer** — `UIApplicationDidReceiveMemoryWarningNotification` as a second line of defense, also cancels generation.
|
| 182 |
+
3. **Pre-flight 500MB check** — `os_proc_available_memory() < 500MB` returns error before starting generation (checked in both `flashmoe_generate` and `flashmoe_generate_continuation`).
|
| 183 |
+
4. **Adaptive context length** — Runtime context cap based on available memory and Metal wired budget at model load time.
|
| 184 |
+
5. **Pre-allocated scratch buffers** — 30+ static scratch buffers allocated once at model load, reused across all layers per token. Eliminates ~300 malloc/free per token.
|
| 185 |
+
6. **Metal buffer nil checks** — All 40+ `newBufferWithLength` calls checked for nil with actionable error messages and early return.
|
| 186 |
+
7. **calloc guards** — All CPU allocations checked for NULL with error reporting.
|
| 187 |
+
8. **posix_memalign for expert I/O** — 2MB-aligned expert data buffers with error checking.
|
| 188 |
+
|
| 189 |
+
See [docs/oom-prevention.md](../docs/oom-prevention.md) for the full architecture document.
|
| 190 |
+
|
| 191 |
+
### ARC Cleanup
|
| 192 |
+
MetalCtx `free()` without nil-ing `id<>` Objective-C fields caused heap corruption on model switch. Fix: nil all `id<>` fields before `free`.
|
| 193 |
+
|
| 194 |
+
### 2-Bit Auto-Detection
|
| 195 |
+
iOS load path was missing 2-bit directory check. Added auto-detection in `flashmoe_load()`.
|
| 196 |
+
|
| 197 |
+
## 397B on iPhone -- What We Tried
|
| 198 |
+
|
| 199 |
+
| Approach | Result | Notes |
|
| 200 |
+
|----------|--------|-------|
|
| 201 |
+
| Metal 4GB buffer workarounds | All failed | See above |
|
| 202 |
+
| K=2 on 397B (trained K=10) | Gibberish | 20% of trained expert capacity |
|
| 203 |
+
| K=4 on 397B (trained K=10) | Degenerate ("!!!!") | 40% capacity insufficient |
|
| 204 |
+
| K=6+ on 397B | Untested | Needs GPU path (split weights) |
|
| 205 |
+
| File Provider Storage | +latency | File coordination overhead on every pread |
|
| 206 |
+
|
| 207 |
+
### Performance Projections (After Split Weights Enable GPU Path)
|
| 208 |
+
|
| 209 |
+
| Configuration | Expert I/O | Expected tok/s |
|
| 210 |
+
|--------------|-----------|----------------|
|
| 211 |
+
| K=10, 4-bit | 4.1 GB/token | ~0.5 |
|
| 212 |
+
| K=4, 4-bit | 1.6 GB/token | ~1.0 |
|
| 213 |
+
| K=4, tiered | 1.1 GB/token | ~1.2 |
|
| 214 |
+
|
| 215 |
+
## GPTQ 2-bit: Path to 397B on iPhone
|
| 216 |
+
|
| 217 |
+
The GPTQ quantization pipeline opens a realistic path to running the full 397B model on 256GB iPhones.
|
| 218 |
+
|
| 219 |
+
**Key insight**: The 20% hot 4-bit + 80% GPTQ 2-bit tiered configuration produces a **134GB** model. This fits on a 256GB iPhone with room for the OS, apps, and user data. Unlike RTN 2-bit (which breaks JSON/tool calling), GPTQ 2-bit uses Hessian-guided error compensation to preserve output quality at 2-bit precision.
|
| 220 |
+
|
| 221 |
+
| Component | Size | Notes |
|
| 222 |
+
|-----------|------|-------|
|
| 223 |
+
| Hot experts (20%, 4-bit) | ~42 GB | Top ~25% by activation frequency, full quality |
|
| 224 |
+
| Cold experts (80%, GPTQ 2-bit) | ~87 GB | GPTQ error compensation fixes JSON output |
|
| 225 |
+
| Non-expert weights | 5.5 GB | Needs split into two <4GB files for Metal, or CPU fallback |
|
| 226 |
+
| **Total** | **~134 GB** | Fits on 256GB iPhone |
|
| 227 |
+
|
| 228 |
+
**Expected quality**: Production-grade JSON and tool calling. GPTQ's column-wise error compensation keeps accumulated quantization error bounded even at 2-bit, avoiding the `\name\` corruption seen with RTN 2-bit.
|
| 229 |
+
|
| 230 |
+
**Remaining requirements**:
|
| 231 |
+
- Split `model_weights.bin` into two <4GB files (Metal per-buffer limit on iOS)
|
| 232 |
+
- Alternatively, use CPU fallback for non-expert weight projections (5.5GB exceeds 4GB Metal limit)
|
| 233 |
+
- Test K=6+ with GPU path on device for coherent output
|
| 234 |
+
- Validate GPTQ 2-bit JSON quality on real tool-calling workloads
|
| 235 |
+
|
| 236 |
+
See [docs/quantization-guide.md](quantization-guide.md) for the full GPTQ pipeline documentation.
|
| 237 |
+
|
| 238 |
+
## Next Steps
|
| 239 |
+
|
| 240 |
+
1. **Split `model_weights.bin` into two <4GB files** -- enables GPU path on iOS for 397B
|
| 241 |
+
2. **Test K=6/8/10 with GPU path** -- find minimum viable K for coherent 397B output
|
| 242 |
+
3. **Upload split 397B model to HuggingFace**
|
| 243 |
+
4. **Adaptive K** -- auto-select based on device RAM and thermal state
|
| 244 |
+
5. **Thermal throttling awareness** -- monitor `ProcessInfo.ThermalState`, reduce K when throttling
|
docs/oom-prevention.md
ADDED
|
@@ -0,0 +1,136 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# OOM Prevention Architecture
|
| 2 |
+
|
| 3 |
+
**Status**: Implemented across all code paths
|
| 4 |
+
**Relevant files**: `metal_ctx.h`, `layer_forward.h`, `generate.h`, `FlashMoEEngine.m`
|
| 5 |
+
|
| 6 |
+
## Problem
|
| 7 |
+
|
| 8 |
+
Flash-MoE runs on devices with as little as 12GB unified memory (iPhone 17) while managing hundreds of Metal buffers, CPU scratch allocations, and multi-gigabyte model weights. A single failed allocation can crash the app or corrupt inference output. On iOS, the system kills apps that exceed their memory budget (jetsam) with no warning beyond memory pressure notifications.
|
| 9 |
+
|
| 10 |
+
## Design Principle: Pre-Allocate, Check Everything, Fail Gracefully
|
| 11 |
+
|
| 12 |
+
The engine follows three rules:
|
| 13 |
+
1. **Pre-allocate at load time** -- all scratch buffers allocated once during `flashmoe_load()`, reused across all layers and tokens. No per-token malloc/free.
|
| 14 |
+
2. **Check every allocation** -- every `newBufferWithLength`, `calloc`, and `posix_memalign` return value is tested. Failures produce actionable error messages and early return (not abort).
|
| 15 |
+
3. **Monitor at runtime** -- iOS memory pressure and warning handlers cancel generation before the system kills the process.
|
| 16 |
+
|
| 17 |
+
## Allocation Hardening
|
| 18 |
+
|
| 19 |
+
### 1. Static Scratch Buffers (layer_forward.h)
|
| 20 |
+
|
| 21 |
+
30+ scratch buffers pre-allocated in `init_layer_scratch()` at model load:
|
| 22 |
+
|
| 23 |
+
```
|
| 24 |
+
s_normed, s_residual, s_attn_proj, s_h_post, s_h_mid,
|
| 25 |
+
s_gate_scores, s_spec_gate_scores, s_shared_gate, s_shared_up,
|
| 26 |
+
s_moe_out, s_shared_out, s_q_proj_out, s_k_proj_out, s_v_proj_out,
|
| 27 |
+
s_q, s_q_gate, s_attn_out, s_qkv_proj_out, s_z_proj_out,
|
| 28 |
+
s_beta_proj_out, s_alpha_proj_out, s_conv_out, s_out_vals, s_gated_out,
|
| 29 |
+
s_expert_out_cpu, s_gate_proj_out, s_up_proj_out, s_act_out,
|
| 30 |
+
s_shared_act, s_k_dequant, s_v_dequant
|
| 31 |
+
```
|
| 32 |
+
|
| 33 |
+
All checked in a single compound `if` statement. If any allocation fails, `init_layer_scratch()` returns -1 and the engine refuses to start inference.
|
| 34 |
+
|
| 35 |
+
**Impact**: Eliminates ~300 malloc/free per token. Each token traverses all layers, and each layer previously allocated temporary buffers for attention projections, expert routing, and MoE combine. Pre-allocation converts heap allocation to pointer reuse.
|
| 36 |
+
|
| 37 |
+
### 2. Metal Buffer Nil Checks (metal_ctx.h)
|
| 38 |
+
|
| 39 |
+
All 40+ Metal buffer allocations in `metal_setup()` are grouped and checked:
|
| 40 |
+
|
| 41 |
+
- Input/output buffers (`buf_input`, `buf_output`)
|
| 42 |
+
- Expert data buffers (`buf_expert_data`, `buf_expert_input`, `buf_expert_gate`, etc.)
|
| 43 |
+
- Multi-expert double-buffered slots (`buf_multi_expert_data[k]`, `buf_multi_expert_data_B[k]`)
|
| 44 |
+
- Shared expert buffers (`buf_shared_gate`, `buf_shared_up`, `buf_shared_act`, `buf_shared_out`)
|
| 45 |
+
- Residual and normalization buffers (`buf_residual`, `buf_h_mid`, `buf_sum_sq`)
|
| 46 |
+
- MoE combine buffers (`buf_moe_hidden`, `buf_combine_params`, `buf_cmd3_sum_sq`)
|
| 47 |
+
- KV cache buffers (`buf_kv_k[i]`, `buf_kv_v[i]` for each layer)
|
| 48 |
+
- FP8 scale buffers (`buf_kv_k_scales[i]`, `buf_kv_v_scales[i]`) when FP8 KV is enabled
|
| 49 |
+
- Attention scratch (`buf_attn_q`, `buf_attn_scores`, `buf_attn_out`, `buf_attn_gate`)
|
| 50 |
+
- Delta-net state (`buf_delta_state[i]`, `buf_conv_state[i]` per layer)
|
| 51 |
+
- Delta-net scratch (`buf_delta_q`, `buf_delta_k`, `buf_delta_v`, etc.)
|
| 52 |
+
|
| 53 |
+
Each check prints the buffer name, requested size, and device name to stderr, then returns early. Metal returns nil when the device cannot satisfy the allocation (typically when wired memory is exhausted).
|
| 54 |
+
|
| 55 |
+
### 3. calloc Guards
|
| 56 |
+
|
| 57 |
+
All CPU heap allocations use `calloc()` (zero-initialized) and check for NULL:
|
| 58 |
+
- Layer file descriptor arrays in `FlashMoEEngine.m`
|
| 59 |
+
- Tracking arrays (`alloc_tracking_arrays()`)
|
| 60 |
+
- Vocabulary and tokenizer data
|
| 61 |
+
|
| 62 |
+
### 4. posix_memalign for Expert I/O (metal_ctx.h)
|
| 63 |
+
|
| 64 |
+
Expert data double-buffers use 2MB-aligned allocation via `posix_memalign()`:
|
| 65 |
+
```c
|
| 66 |
+
int pa_ret1 = posix_memalign(&aligned_data, 2*1024*1024, expert_alloc_size);
|
| 67 |
+
int pa_ret2 = posix_memalign(&aligned_data_b, 2*1024*1024, expert_alloc_size);
|
| 68 |
+
```
|
| 69 |
+
2MB alignment matches macOS huge page size, enabling the OS to use huge pages for expert I/O buffers. The return value is checked and logged.
|
| 70 |
+
|
| 71 |
+
## iOS Memory Management
|
| 72 |
+
|
| 73 |
+
### 5. Memory Pressure Dispatch Source (FlashMoEEngine.m)
|
| 74 |
+
|
| 75 |
+
A GCD dispatch source monitors system memory pressure:
|
| 76 |
+
```
|
| 77 |
+
DISPATCH_SOURCE_TYPE_MEMORYPRESSURE
|
| 78 |
+
DISPATCH_MEMORYPRESSURE_WARN | DISPATCH_MEMORYPRESSURE_CRITICAL
|
| 79 |
+
```
|
| 80 |
+
|
| 81 |
+
On `DISPATCH_MEMORYPRESSURE_CRITICAL`: sets `atomic_store(&ctx->cancelled, 1)`, which causes the generation loop to exit at the next token boundary. This is an **actionable** handler -- it stops generation, not just logs.
|
| 82 |
+
|
| 83 |
+
On `DISPATCH_MEMORYPRESSURE_WARN`: logs a warning (generation continues but the system is under pressure).
|
| 84 |
+
|
| 85 |
+
### 6. didReceiveMemoryWarning Observer (FlashMoEEngine.m)
|
| 86 |
+
|
| 87 |
+
A second line of defense via `UIApplicationDidReceiveMemoryWarningNotification`:
|
| 88 |
+
```objc
|
| 89 |
+
[[NSNotificationCenter defaultCenter]
|
| 90 |
+
addObserverForName:UIApplicationDidReceiveMemoryWarningNotification ...]
|
| 91 |
+
```
|
| 92 |
+
Also cancels generation via `atomic_store(&ctx->cancelled, 1)`. This fires when UIKit receives a memory warning from the system, which may arrive on a different schedule than the dispatch source.
|
| 93 |
+
|
| 94 |
+
Both handlers are registered at model load and removed at unload.
|
| 95 |
+
|
| 96 |
+
### 7. Pre-Flight 500MB Check (FlashMoEEngine.m)
|
| 97 |
+
|
| 98 |
+
Before starting any generation (both `flashmoe_generate` and `flashmoe_generate_continuation`):
|
| 99 |
+
```c
|
| 100 |
+
size_t avail = os_proc_available_memory();
|
| 101 |
+
if (avail < 500 * 1024 * 1024) {
|
| 102 |
+
snprintf(ctx->last_error, ..., "Insufficient memory (%.0f MB available, need 500+ MB)");
|
| 103 |
+
return -1;
|
| 104 |
+
}
|
| 105 |
+
```
|
| 106 |
+
|
| 107 |
+
**Why 500MB**: Generation needs memory for expert I/O buffers (~27MB per token for K=4 at 4-bit), Metal command buffer overhead, and headroom for the OS page cache to service expert reads. 500MB provides a safety margin that prevents starting a generation that would immediately trigger memory pressure.
|
| 108 |
+
|
| 109 |
+
### 8. Adaptive Context Length (FlashMoEEngine.m)
|
| 110 |
+
|
| 111 |
+
At model load, `os_proc_available_memory()` determines the maximum KV cache size:
|
| 112 |
+
```
|
| 113 |
+
available_memory -> bytes_per_position (float32 or FP8) -> max_context
|
| 114 |
+
```
|
| 115 |
+
With FP8 KV cache enabled, `bytes_per_elem = 1` instead of `sizeof(float) = 4`, enabling 4x longer context within the same memory budget.
|
| 116 |
+
|
| 117 |
+
Additionally, `recommendedMaxWorkingSetSize` from the Metal device constrains the total GPU buffer allocation to stay within the wired memory budget.
|
| 118 |
+
|
| 119 |
+
## Prompt Length Caps
|
| 120 |
+
|
| 121 |
+
All prompt input paths cap token count to `cfg.max_seq_len`:
|
| 122 |
+
- System prompt tokenization (chat mode)
|
| 123 |
+
- User prompt tokenization (single-shot mode)
|
| 124 |
+
- HTTP API prompt tokenization (serve mode)
|
| 125 |
+
|
| 126 |
+
Excess tokens are silently truncated with a log message, preventing KV cache overflow.
|
| 127 |
+
|
| 128 |
+
## Architecture Decision: Why Pre-Allocate
|
| 129 |
+
|
| 130 |
+
The alternative -- allocating per-token -- was the original approach and caused two problems:
|
| 131 |
+
|
| 132 |
+
1. **Fragmentation**: 300 malloc/free per token over thousands of tokens fragments the heap. On iOS with 12GB, fragmentation can cause allocation failures even when total memory is available.
|
| 133 |
+
|
| 134 |
+
2. **Latency variance**: malloc can trigger page faults, madvise calls, or zone coalescing at unpredictable times. Pre-allocation moves all of this to model load, where latency is acceptable.
|
| 135 |
+
|
| 136 |
+
The pre-allocation cost is modest: ~2-3MB of scratch buffers for the 397B model (hidden_dim=4096, moe_intermediate=24576). This is negligible compared to the 200MB+ of Metal buffers.
|
docs/optimization-experiments-q4.md
ADDED
|
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Q4 Expert Optimization Experiments
|
| 2 |
+
|
| 3 |
+
## Context
|
| 4 |
+
|
| 5 |
+
After discovering that 2-bit expert quantization broke tool calling (JSON quotes → backslashes), we reverted to 4-bit experts. This dropped performance from 5.74 tok/s (2-bit) to 3.50 tok/s (4-bit). The goal: recover as much speed as possible while maintaining 4-bit quality.
|
| 6 |
+
|
| 7 |
+
The 4-bit experts are 7,077,888 bytes each (6.75 MB). With K=4 active experts per layer and 60 layers, each token reads 240 experts = 1.68 GB from SSD.
|
| 8 |
+
|
| 9 |
+
## Baseline Pipeline (4-bit, K=4, trust OS page cache)
|
| 10 |
+
|
| 11 |
+
```
|
| 12 |
+
Per layer (4.28 ms avg):
|
| 13 |
+
cmd1_wait: 1.22 ms (28%) GPU: CMD3(prev) + CMD1 attention projections
|
| 14 |
+
cmd2_wait: 0.55 ms (13%) GPU: o_proj + norm + routing + shared expert
|
| 15 |
+
expert_io: 2.41 ms (56%) SSD: 4×7MB parallel pread
|
| 16 |
+
CPU work: 0.10 ms ( 2%) encode + attention + routing + memcpy
|
| 17 |
+
|
| 18 |
+
60 layers × 4.28 ms = 257 ms per token = 3.90 tok/s
|
| 19 |
+
```
|
| 20 |
+
|
| 21 |
+
Page cache hit rate: ~71% (35 GB cache, 209 GB model).
|
| 22 |
+
Warm cache parallel pread: 1.0 ms. Cold SSD: 5.8 ms. Mixed: 2.4 ms.
|
| 23 |
+
|
| 24 |
+
## Experiment Results
|
| 25 |
+
|
| 26 |
+
### Kept: FMA Dequant Kernel (+2.6% → 4.36 tok/s)
|
| 27 |
+
|
| 28 |
+
Rearranged the inner loop of `dequant_matvec_4bit_v3` from:
|
| 29 |
+
```metal
|
| 30 |
+
acc += (float(nibble) * scale + bias) * x;
|
| 31 |
+
```
|
| 32 |
+
to:
|
| 33 |
+
```metal
|
| 34 |
+
float sx = scale * x, bx = bias * x;
|
| 35 |
+
acc += fma(float(nibble), sx, bx);
|
| 36 |
+
```
|
| 37 |
+
|
| 38 |
+
Pre-computing `scale*x` and `bias*x` per input element allows the GPU to use the fused multiply-add unit for the dequant+multiply in one instruction. Reduces per-nibble cost from (convert + mul + add + mul + add) to (convert + fma + add).
|
| 39 |
+
|
| 40 |
+
Impact: cmd1_wait -5.4%, cmd2_wait -10.7%. Total: 3.90 → 4.36 tok/s.
|
| 41 |
+
|
| 42 |
+
### Discarded: LZ4 Expert Compression (-13%)
|
| 43 |
+
|
| 44 |
+
Repacked 209 GB of expert files to 175 GB with LZ4 compression. Apple's LZ4 decompressor runs at 41 GB/s (NEON hardware-accelerated), making decompression only 0.17 ms per expert.
|
| 45 |
+
|
| 46 |
+
Results:
|
| 47 |
+
- Isolated cold reads: 15-24% faster (less data from SSD)
|
| 48 |
+
- Isolated decompression: 0.17 ms at 41 GB/s (essentially free)
|
| 49 |
+
- **Full pipeline: 3.55 tok/s (-13%)** — the 0.68 ms/layer decompress cost exceeds the warm cache I/O savings. The OS page cache is efficient enough that most reads are warm.
|
| 50 |
+
|
| 51 |
+
Also tested LZFSE (2.6 GB/s, too slow), APFS transparent compression (kernel serializes read+decompress, 2× slower), and per-expert files (15% slower from VFS metadata overhead).
|
| 52 |
+
|
| 53 |
+
Key finding: Apple's M3 Max SSD is so fast that CPU-based decompression can't keep up for warm cache reads. LZ4 only wins for cold reads, but the page cache handles most reads.
|
| 54 |
+
|
| 55 |
+
### Discarded: Expert Routing Prediction (-18%)
|
| 56 |
+
|
| 57 |
+
Built a temporal prediction system: store previous token's expert routing per layer, prefetch those experts into double-buffered Metal buffers during the next token's CMD1 wait.
|
| 58 |
+
|
| 59 |
+
Results:
|
| 60 |
+
- Temporal hit rate: 25.6% (only 1 of 4 experts matches between tokens)
|
| 61 |
+
- The 75% misses waste SSD bandwidth and require sync pread after the prediction wait
|
| 62 |
+
- With K=4 parallel reads, wall time = max(4 reads). Need ALL 4 to hit for improvement.
|
| 63 |
+
- P(all 4 hit) at 25% = 0.25⁴ = 0.4%. Practically zero.
|
| 64 |
+
|
| 65 |
+
Also trained an MLP predictor (31% accuracy from pre-attention hidden state — worse than temporal baseline). The gate_proj "logit lens" approach achieves 53% from pre-attention state, but the K=4 exponential penalty still kills it.
|
| 66 |
+
|
| 67 |
+
### Discarded: F_RDADVISE Prefetch (net 0%)
|
| 68 |
+
|
| 69 |
+
Sent F_RDADVISE kernel hints between CMD1 commit and wait to prefetch next token's predicted experts during GPU compute.
|
| 70 |
+
|
| 71 |
+
Results:
|
| 72 |
+
- expert_io: -31% (page cache warming works!)
|
| 73 |
+
- cmd2_wait: +73% (GPU memory bandwidth contention from SSD DMA)
|
| 74 |
+
- **Net: 0% across 5 diverse prompts**
|
| 75 |
+
|
| 76 |
+
Root cause: Apple Silicon unified memory architecture. SSD DMA and GPU matvec share the same memory controller. The GPU's dequant kernels are bandwidth-saturated at 418 GiB/s. Even 17.5 GB/s of background DMA (~4%) causes disproportionate latency spikes through memory controller arbitration. This is architectural — cannot be worked around in software.
|
| 77 |
+
|
| 78 |
+
### Discarded: GPU Kernel Variants
|
| 79 |
+
|
| 80 |
+
- **LUT dequant (v5)**: Pre-compute 16-entry lookup table per group to eliminate uint→float conversions. -2% because GPU indirect register access serializes.
|
| 81 |
+
- **Vector load (v4)**: uint4 loads for coalesced memory access. -3% from register pressure.
|
| 82 |
+
- **extract_bits intrinsic**: Neutral — compiler already generates the same instruction.
|
| 83 |
+
- **Spin-poll GPU wait**: -23%. CPU spinning steals thermal budget from GPU on unified architecture.
|
| 84 |
+
- **addCompletedHandler**: Neutral in practice — isolated 20% win on micro-benchmark but real workloads have enough GPU compute to hide the wait overhead.
|
| 85 |
+
|
| 86 |
+
### Discarded: I/O Path Alternatives
|
| 87 |
+
|
| 88 |
+
- **dispatch_io**: -70%. Apple's GCD I/O framework adds dispatch_data management overhead (allocate, map, memcpy) that far exceeds any kernel scheduling benefit.
|
| 89 |
+
- **aio_read**: -7% (matches GCD group + pread, which we already use).
|
| 90 |
+
- **Expert file clustering**: 0%. NVMe doesn't care about scatter distance at 7MB read granularity. 4 reads spanning 21 MB vs 2.9 GB take the same time.
|
| 91 |
+
- **GPU private buffer compression**: Isolated -13.5% per matvec (GPU hardware memory compression on StorageModePrivate). But in pipeline: blitting 4×7MB shared→private costs more than the matvec savings. -20% overall.
|
| 92 |
+
|
| 93 |
+
### Analyzed but not implemented: MTP Speculative Decoding
|
| 94 |
+
|
| 95 |
+
Qwen 3.5 ships with an MTP (Multi-Token Prediction) head — a single MoE transformer layer that predicts the next-next token. The head exists in the model config (`mtp_num_hidden_layers: 1`) but weights were stripped from the MLX quantization.
|
| 96 |
+
|
| 97 |
+
Analysis showed MTP speculative decoding doesn't help for MoE with SSD streaming: each speculated token requires its OWN expert routing and I/O. Batched verification of 2 tokens costs ~1.75× expert I/O for 1.7 tokens (70% acceptance). Break-even at best.
|
| 98 |
+
|
| 99 |
+
This contrasts with dense models where verification cost is constant regardless of batch size (same weights for every token).
|
| 100 |
+
|
| 101 |
+
## The Unified Memory Constraint
|
| 102 |
+
|
| 103 |
+
The single most important finding: **on Apple Silicon, SSD DMA and GPU compute cannot be profitably overlapped.** They share the same memory controller, and the GPU's dequant kernels are bandwidth-saturated. Any background I/O during GPU compute causes disproportionate GPU slowdown.
|
| 104 |
+
|
| 105 |
+
This means the serial pipeline (GPU → SSD → GPU) is actually **hardware-optimal** for this architecture. The current pipeline already achieves the best possible scheduling.
|
| 106 |
+
|
| 107 |
+
## Summary
|
| 108 |
+
|
| 109 |
+
| Configuration | tok/s | Status |
|
| 110 |
+
|--------------|-------|--------|
|
| 111 |
+
| 2-bit experts (best speed) | 5.74 | Quality regression (broken JSON) |
|
| 112 |
+
| 2-bit peak single token | 7.05 | Warm cache burst |
|
| 113 |
+
| **4-bit + FMA kernel** | **4.36** | **Current best. Quality preserved.** |
|
| 114 |
+
| 4-bit baseline (no FMA) | 3.90 | Previous 4-bit baseline |
|
| 115 |
+
| 4-bit + LZ4 compression | 3.55 | Decompress overhead > I/O savings |
|
| 116 |
+
| 4-bit + temporal prediction | 3.18 | 25% hit rate wastes SSD bandwidth |
|
| 117 |
+
| 4-bit + F_RDADVISE prefetch | 3.91 | GPU contention cancels I/O savings |
|
| 118 |
+
|
| 119 |
+
The 4-bit performance ceiling on M3 Max 48GB is approximately **4.4 tok/s** for sustained generation, limited by:
|
| 120 |
+
- 56% SSD expert I/O (2.4 ms/layer, hardware-limited)
|
| 121 |
+
- 41% GPU dequant matvec (1.8 ms/layer, bandwidth-limited)
|
| 122 |
+
- 3% CPU overhead (0.1 ms/layer)
|
| 123 |
+
|
| 124 |
+
Further improvement requires either hardware changes (more RAM for expert caching, faster SSD) or model architecture changes (fewer/smaller experts, larger shared expert to reduce per-token I/O).
|
| 125 |
+
|
| 126 |
+
## iPhone Testing (A19, 12GB)
|
| 127 |
+
|
| 128 |
+
The same optimization techniques were tested on iPhone 17 with the 35B model:
|
| 129 |
+
|
| 130 |
+
### Applicable Optimizations
|
| 131 |
+
|
| 132 |
+
| Optimization | Result on iPhone | Notes |
|
| 133 |
+
|-------------|-----------------|-------|
|
| 134 |
+
| FMA dequant kernel | Included in baseline | Same kernel runs on both platforms |
|
| 135 |
+
| CMD1+CMD2 merge | +5-10% projected | Saves 30 sync points per token on 35B (30 linear layers) |
|
| 136 |
+
| Fused expert kernel | +3-5% projected | Reduces dispatch overhead per expert |
|
| 137 |
+
| Expert prefetch | Under validation | Depends on A19 NVMe characteristics |
|
| 138 |
+
| FP16 accumulation | Under validation | A19 has fp16 ALUs, may be more impactful than M3 Max |
|
| 139 |
+
| K-reduction (K=4 from K=8) | **~2x speedup** | Halves expert I/O, the dominant bottleneck |
|
| 140 |
+
|
| 141 |
+
### iPhone-Specific Results
|
| 142 |
+
|
| 143 |
+
| Configuration | tok/s | Notes |
|
| 144 |
+
|--------------|-------|-------|
|
| 145 |
+
| 35B, K=8, baseline | 5.5 | Full quality, 19.5GB model |
|
| 146 |
+
| 35B, K=8, tiered | 5.5+ | 13.4GB model, same quality |
|
| 147 |
+
| 35B, K=4, projected with all opts | ~11 | CMD merge + fused expert + K reduction |
|
| 148 |
+
| 397B, K=4, CPU fallback | ~0.003 | Metal 4GB buffer limit blocks GPU path |
|
| 149 |
+
|
| 150 |
+
### Key Differences from MacBook
|
| 151 |
+
|
| 152 |
+
1. **NVMe bandwidth** -- iPhone ~2.5-3 GB/s vs MacBook 17.5 GB/s. Expert I/O takes proportionally longer, making K-reduction even more impactful on mobile.
|
| 153 |
+
2. **Memory budget** -- 12GB vs 48GB. Less room for OS page cache, so expert cache hit rates are lower. FP8 KV cache and sliding window attention free up memory for the page cache.
|
| 154 |
+
3. **Thermal constraints** -- iPhone thermal throttles sooner. Sustained generation speeds may be lower than burst measurements.
|
| 155 |
+
4. **Metal buffer limits** -- 4GB per buffer on iOS vs unlimited on macOS. Prevents GPU path for 397B model without split weight files.
|
docs/plan-async-pread-pipeline.md
ADDED
|
@@ -0,0 +1,166 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Plan: Cross-Layer Async Pread Pipeline
|
| 2 |
+
|
| 3 |
+
## Status
|
| 4 |
+
- Async pread mechanism implemented and working (async_pread_start/wait in infer.m)
|
| 5 |
+
- Within-layer overlap tested: no improvement (only 0.1ms overlap window)
|
| 6 |
+
- Need: CROSS-LAYER overlap for ~2ms of pread hiding
|
| 7 |
+
|
| 8 |
+
## Current Per-Layer Sequence (4.5ms total)
|
| 9 |
+
```
|
| 10 |
+
[deferred_wait] → [CMD1 submit+wait] → [CPU attn] → [CMD2 submit+wait] → [routing] → [SYNC pread] → [CMD3 submit]
|
| 11 |
+
0.87ms 0.5ms 0.27ms 0.45ms 0.003ms 2.43ms 0.03ms
|
| 12 |
+
```
|
| 13 |
+
|
| 14 |
+
## Target Sequence
|
| 15 |
+
```
|
| 16 |
+
Layer N: ... → [routing] → [START async pread into BUF_A] → [CMD3 submit (using BUF_B from prev)]
|
| 17 |
+
Layer N+1: [deferred_wait] → [CMD1] → [CPU attn] → [CMD2] → [routing] → [WAIT async pread BUF_A] → [CMD3 submit (using BUF_A)]
|
| 18 |
+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
| 19 |
+
~2.1ms of compute overlapping with N's pread
|
| 20 |
+
```
|
| 21 |
+
|
| 22 |
+
Pread for layer N runs during layer N+1's compute. By the time N+1 needs expert data, N's pread has had 2.1ms of head start.
|
| 23 |
+
|
| 24 |
+
## Implementation Steps
|
| 25 |
+
|
| 26 |
+
### 1. Double-Buffer Expert Data
|
| 27 |
+
Already have: `buf_multi_expert_data[MAX_K]` (set A) and `buf_multi_expert_data_B[MAX_K]` (set B).
|
| 28 |
+
|
| 29 |
+
Add a flip flag:
|
| 30 |
+
```c
|
| 31 |
+
static int g_expert_buf_flip = 0; // 0 = use set A for current, 1 = use set B
|
| 32 |
+
```
|
| 33 |
+
|
| 34 |
+
Each layer alternates which buffer set it writes pread data INTO vs which set CMD3 reads FROM.
|
| 35 |
+
|
| 36 |
+
### 2. Restructure fused_layer_forward
|
| 37 |
+
|
| 38 |
+
**At layer start (after deferred_wait):**
|
| 39 |
+
- If async pread is in flight from previous layer: DON'T wait yet
|
| 40 |
+
- Continue with CMD1, CPU attn, CMD2, routing
|
| 41 |
+
|
| 42 |
+
**After routing:**
|
| 43 |
+
- NOW wait for the previous layer's async pread (it's had ~2ms to complete)
|
| 44 |
+
- Start THIS layer's async pread into the OTHER buffer set
|
| 45 |
+
- Encode CMD3 using the COMPLETED buffer set (from previous layer's pread)
|
| 46 |
+
|
| 47 |
+
Wait — this doesn't work because CMD3 needs THIS layer's expert data, not the previous layer's. Let me rethink.
|
| 48 |
+
|
| 49 |
+
### Correct Design
|
| 50 |
+
|
| 51 |
+
The pread for layer N needs to complete before layer N's CMD3 encodes. But we want pread for N to overlap with layer N's CMD1+attn+CMD2 (which don't use expert data).
|
| 52 |
+
|
| 53 |
+
**Revised flow:**
|
| 54 |
+
```
|
| 55 |
+
Layer N start:
|
| 56 |
+
1. Wait for N-1's deferred CMD3 (or GPU combine)
|
| 57 |
+
2. Submit CMD1 (attention projections)
|
| 58 |
+
3. [BACKGROUND: start async pread for layer N's experts]
|
| 59 |
+
- But we don't know N's experts yet! Routing hasn't happened.
|
| 60 |
+
```
|
| 61 |
+
|
| 62 |
+
This is the fundamental problem: we can't start pread until after routing, but routing is the LAST thing before pread in the pipeline.
|
| 63 |
+
|
| 64 |
+
### The Real Solution: Decouple Routing from Expert Loading
|
| 65 |
+
|
| 66 |
+
Split CMD2 into two parts:
|
| 67 |
+
- CMD2a: o_proj + residual + norm (produces h_post for routing)
|
| 68 |
+
- CMD2b: routing gate_proj (produces gate_scores)
|
| 69 |
+
|
| 70 |
+
Then:
|
| 71 |
+
```
|
| 72 |
+
Layer N:
|
| 73 |
+
CMD1 → CPU attn → CMD2a+CMD2b → wait → routing topK → [START async pread] → CMD3 (deferred)
|
| 74 |
+
Layer N+1:
|
| 75 |
+
[async pread from N still running]
|
| 76 |
+
deferred_wait → CMD1 → CPU attn → CMD2a+CMD2b → wait → routing topK
|
| 77 |
+
[NOW wait for N's async pread — it's had the entire N+1 compute time]
|
| 78 |
+
→ CMD3 using N+1's expert data that we NOW start loading synchronously
|
| 79 |
+
```
|
| 80 |
+
|
| 81 |
+
Hmm, this still doesn't help because we need N+1's experts, not N's.
|
| 82 |
+
|
| 83 |
+
### Actually Correct Solution: Pipeline Expert Data One Layer Ahead
|
| 84 |
+
|
| 85 |
+
The insight: at the end of layer N, we have N's expert data loaded. We submit CMD3 (deferred) which uses that data. CMD3 runs on GPU while we start layer N+1.
|
| 86 |
+
|
| 87 |
+
If we started loading N+1's experts AT THE SAME TIME as submitting N's CMD3:
|
| 88 |
+
```
|
| 89 |
+
Layer N end: [submit CMD3_N using BUF_A] + [start async pread for N+1 into BUF_B]
|
| 90 |
+
Layer N+1: [deferred_wait N] → [CMD1] → [attn] → [CMD2] → [routing]
|
| 91 |
+
[async pread N+1 completes during this time]
|
| 92 |
+
→ [check: do loaded experts match routing? if yes, use BUF_B; if no, sync pread]
|
| 93 |
+
```
|
| 94 |
+
|
| 95 |
+
But we DON'T KNOW layer N+1's experts at the end of layer N. We'd need to PREDICT them.
|
| 96 |
+
|
| 97 |
+
Previous prediction attempts failed (53% accuracy, overhead > benefit).
|
| 98 |
+
|
| 99 |
+
### ALTERNATIVE: Overlap pread with CMD3 GPU execution
|
| 100 |
+
|
| 101 |
+
Currently CMD3 is deferred — GPU runs it while we start the next layer. But we DON'T start loading the next layer's experts during CMD3. What if we did?
|
| 102 |
+
|
| 103 |
+
After CMD3 submit for layer N:
|
| 104 |
+
```
|
| 105 |
+
[submit CMD3_N] → [start next layer's CMD1+attn+CMD2] → [routing N+1] → [pread N+1]
|
| 106 |
+
↑ CMD3_N runs on GPU here, overlapping with N+1's compute
|
| 107 |
+
```
|
| 108 |
+
|
| 109 |
+
The pread for N+1 currently starts AFTER routing for N+1, which is after CMD2 for N+1, which is after deferred_wait for CMD3_N. So the pread can't start until CMD3_N is done.
|
| 110 |
+
|
| 111 |
+
But with GPU combine+norm in CMD3, we eliminated the deferred_wait. CMD1 for N+1 submits immediately after CMD3_N. The GPU executes CMD3_N → CMD1_N+1 back-to-back. The CPU is free during this time to do... nothing useful, because it's waiting for CMD1_N+1 to complete.
|
| 112 |
+
|
| 113 |
+
### THE REAL REAL SOLUTION: Start pread during CMD1 wait
|
| 114 |
+
|
| 115 |
+
CMD1_wait takes 0.87ms (includes CMD3_prev + CMD1 GPU time). During that 0.87ms, the CPU is IDLE waiting for GPU. What if the CPU started the pread during that wait?
|
| 116 |
+
|
| 117 |
+
But we don't have the routing results yet — routing happens after CMD2.
|
| 118 |
+
|
| 119 |
+
UNLESS we use the PREVIOUS TOKEN's routing for the same layer as a prediction. This is temporal locality — 20-35% overlap between tokens at the same layer.
|
| 120 |
+
|
| 121 |
+
We already have the prediction infrastructure (`g_prefetch_experts`). The issue was that F_RDADVISE predictions wasted SSD bandwidth. But what about LOADING into actual buffers?
|
| 122 |
+
|
| 123 |
+
```
|
| 124 |
+
Layer N, token T:
|
| 125 |
+
[CMD1 submit] → [while waiting: pread PREDICTED experts into BUF_B based on token T-1]
|
| 126 |
+
→ [CMD1 wait returns] → [CPU attn] → [CMD2] → [routing]
|
| 127 |
+
→ [check predictions: how many of K=4 match?]
|
| 128 |
+
→ [pread only the MISSES into BUF_A (typically 2-3 instead of 4)]
|
| 129 |
+
→ [CMD3 using mix of BUF_A (misses) and BUF_B (hits)]
|
| 130 |
+
```
|
| 131 |
+
|
| 132 |
+
With 30% hit rate: 1.2 of 4 experts are pre-loaded. Saves ~30% of pread time.
|
| 133 |
+
With 50% hit rate: 2 of 4 pre-loaded. Saves ~50%.
|
| 134 |
+
|
| 135 |
+
The difference from before: we're not polluting any CACHE. We're loading into scratch buffers. Predictions that miss are just overwritten. No cache eviction.
|
| 136 |
+
|
| 137 |
+
The cost: gate_proj matvec (~0.1ms) + predicted pread during CMD1_wait (runs during idle CPU time, so ~0ms additional). Net cost is ~0.1ms per layer.
|
| 138 |
+
|
| 139 |
+
## Key Files
|
| 140 |
+
- `infer.m`: `fused_layer_forward()` around line 4900-5230
|
| 141 |
+
- `async_pread_start/wait` at line ~3011-3050
|
| 142 |
+
- `g_prefetch_experts` at line ~195 (temporal prediction state)
|
| 143 |
+
- `buf_multi_expert_data` (set A) and `buf_multi_expert_data_B` (set B) in MetalCtx
|
| 144 |
+
|
| 145 |
+
## Baseline
|
| 146 |
+
- 4-bit, K=4, Trust OS, no cache: 3.50-3.70 tok/s
|
| 147 |
+
- expert_io: 2.43ms/layer
|
| 148 |
+
- Target: reduce to ~1.5ms/layer → ~4.5-5.0 tok/s
|
| 149 |
+
|
| 150 |
+
## Risk
|
| 151 |
+
- SSD bandwidth contention between predicted and actual preads
|
| 152 |
+
- Previous speculative attempts all failed or were neutral
|
| 153 |
+
- Double-buffer complexity
|
| 154 |
+
- Must verify quality is preserved (same output with/without optimization)
|
| 155 |
+
|
| 156 |
+
## Previous Attempts (ALL FAILED)
|
| 157 |
+
- Speculative early routing on pre-attention state: 53% accuracy, cache pollution → slower
|
| 158 |
+
- F_RDADVISE hints: NVMe command contention → slower
|
| 159 |
+
- Temporal F_RDADVISE with lead time: 65-80% wrong predictions → slower
|
| 160 |
+
- mmap memcpy: 5.5x slower for cold data (page faults)
|
| 161 |
+
|
| 162 |
+
## What's Different This Time
|
| 163 |
+
- Loading into SCRATCH buffers (no cache pollution)
|
| 164 |
+
- Using CMD1_wait idle time (no additional CPU cost)
|
| 165 |
+
- Only predicting from previous token at same layer (simple, no gate_proj overhead if we just reuse stored indices)
|
| 166 |
+
- Only need to sync-pread the MISSES, not all 4 experts
|
docs/plan-io-experiments.md
ADDED
|
@@ -0,0 +1,198 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Plan: I/O Optimization Experiments
|
| 2 |
+
|
| 3 |
+
## Baseline Reference
|
| 4 |
+
- Per-layer: 4.28ms total, 2.41ms expert_io (56%)
|
| 5 |
+
- Expert read: 4 × 7MB parallel pread from 3.4GB file
|
| 6 |
+
- Measured: 5.8ms cold parallel, 1.0ms warm parallel, 2.4ms mixed (71% cache hit)
|
| 7 |
+
- Theoretical floor: 28MB / 17.5 GB/s = 1.6ms
|
| 8 |
+
- Gap: 0.8ms overhead per layer (kernel VFS + page cache + NVMe scheduling)
|
| 9 |
+
|
| 10 |
+
---
|
| 11 |
+
|
| 12 |
+
## Experiment 1: dispatch_io vs pread
|
| 13 |
+
|
| 14 |
+
### Isolated test
|
| 15 |
+
Read 4 experts (7MB each) from a layer file using:
|
| 16 |
+
- (A) 4 × pread on 4 pthreads (current approach)
|
| 17 |
+
- (B) 4 × dispatch_io_read on a DISPATCH_IO_RANDOM channel
|
| 18 |
+
- Both with F_NOCACHE to force SSD reads. Both with warm cache. 50 iterations each.
|
| 19 |
+
|
| 20 |
+
Measure: wall time, throughput (GB/s logical).
|
| 21 |
+
|
| 22 |
+
### What dispatch_io does differently
|
| 23 |
+
- Creates a kernel-side I/O channel with optimized scheduling
|
| 24 |
+
- The kernel sees all reads as part of one channel → can reorder NVMe commands by LBA
|
| 25 |
+
- Automatic cleanup handlers (no thread join overhead)
|
| 26 |
+
- May use a different VFS code path optimized for random access
|
| 27 |
+
|
| 28 |
+
### Pipeline contention analysis
|
| 29 |
+
- dispatch_io is async with completion blocks on GCD queues
|
| 30 |
+
- The completion block runs on a GCD thread (same as our async_pread)
|
| 31 |
+
- Memory path: SSD → DMA → DRAM → Metal shared buffer (same as pread)
|
| 32 |
+
- **No new contention** — same memory path as current approach
|
| 33 |
+
- Risk: dispatch_io might add GCD overhead that exceeds VFS savings
|
| 34 |
+
- Risk: completion blocks might have higher latency than pthread wakeup
|
| 35 |
+
|
| 36 |
+
### Expected impact
|
| 37 |
+
- Best case: 10-20% expert_io reduction (eliminates per-syscall VFS overhead)
|
| 38 |
+
- Worst case: neutral or slight regression (GCD completion overhead)
|
| 39 |
+
- Pipeline impact: drop-in replacement for pread, no GPU interaction
|
| 40 |
+
|
| 41 |
+
---
|
| 42 |
+
|
| 43 |
+
## Experiment 2: GPU private buffer compression
|
| 44 |
+
|
| 45 |
+
### Isolated test
|
| 46 |
+
- (A) GPU matvec reading from StorageModeShared buffer (current)
|
| 47 |
+
- (B) GPU blit shared→private, then matvec reading from StorageModePrivate buffer
|
| 48 |
+
- Use same expert weight data, same kernel, same dimensions
|
| 49 |
+
- Measure: blit time, matvec time, total. 100 iterations.
|
| 50 |
+
|
| 51 |
+
### What GPU compression does
|
| 52 |
+
- StorageModePrivate buffers live in GPU-managed memory
|
| 53 |
+
- The GPU's memory controller can apply lossless compression (similar to console GPU
|
| 54 |
+
texture compression — transparent to shaders)
|
| 55 |
+
- For compressible data: effective bandwidth doubles (read 64B, decompress to 128B)
|
| 56 |
+
- 4-bit quantized weights with 2.4-3.7 bits entropy → highly compressible
|
| 57 |
+
- The shader code doesn't change at all — compression is hardware-transparent
|
| 58 |
+
|
| 59 |
+
### Pipeline contention analysis
|
| 60 |
+
- Blit (shared→private) runs on the GPU command queue
|
| 61 |
+
- It would go BEFORE the matvec dispatches in CMD3
|
| 62 |
+
- Timeline: [pread→shared buf] → [GPU blit 0.02ms] → [GPU matvec from private]
|
| 63 |
+
- The blit adds ~0.02ms per expert to CMD3
|
| 64 |
+
- But the matvec might be 30-50% faster from doubled bandwidth
|
| 65 |
+
- **Key contention**: the blit and matvec are both on the same GPU queue (serial)
|
| 66 |
+
The blit cannot overlap with the matvec. It's purely: does the bandwidth gain
|
| 67 |
+
from compression exceed the blit cost?
|
| 68 |
+
- **Memory**: private buffers use GPU-managed memory. 4 × 7MB = 28MB of private
|
| 69 |
+
memory per layer. The GPU manages this pool — may cause memory pressure if the
|
| 70 |
+
pool grows. Need to reuse/recycle the private buffers each layer.
|
| 71 |
+
|
| 72 |
+
### Expected impact
|
| 73 |
+
- Best case: 15-30% cmd1_wait reduction (CMD3 expert matvec faster)
|
| 74 |
+
- Worst case: slight regression (blit cost exceeds compression benefit)
|
| 75 |
+
- Pipeline impact: affects GPU phases only, no SSD interaction
|
| 76 |
+
|
| 77 |
+
---
|
| 78 |
+
|
| 79 |
+
## Experiment 3: Expert file clustering by co-occurrence
|
| 80 |
+
|
| 81 |
+
### Isolated test
|
| 82 |
+
- Run 500 tokens with --freq, collect per-layer expert co-occurrence matrix
|
| 83 |
+
- For each layer: cluster the 512 experts so frequently co-occurring experts are adjacent
|
| 84 |
+
- Repack each layer file with the new ordering (+ save the permutation map)
|
| 85 |
+
- Measure: 4-expert parallel pread with original vs clustered ordering
|
| 86 |
+
- Use same expert indices (mapped through permutation), F_NOCACHE, 50 iterations
|
| 87 |
+
|
| 88 |
+
### What clustering does
|
| 89 |
+
- NVMe SSDs read in pages (4KB-16KB). When we read expert 37 (7MB at offset 262MB),
|
| 90 |
+
the SSD reads pages 262.0-269.0 MB. Expert 38 is at 269-276 MB — adjacent.
|
| 91 |
+
- If the routing selects experts {37, 42, 100, 205}, those are at offsets
|
| 92 |
+
{262, 297, 708, 1451} MB — widely scattered
|
| 93 |
+
- If we reorder so co-occurring experts are adjacent: {37, 42, 100, 205} might become
|
| 94 |
+
physical positions {0, 1, 2, 3} — a 28MB sequential read instead of 4 scattered reads
|
| 95 |
+
- Sequential 28MB at 17.5 GB/s = 1.6ms vs scattered 4×7MB at ~5.8ms
|
| 96 |
+
|
| 97 |
+
### Pipeline contention analysis
|
| 98 |
+
- This changes the FILE LAYOUT only — the inference code reads the same way
|
| 99 |
+
- Expert indices get mapped through a permutation table (one array lookup, ~0ns)
|
| 100 |
+
- **No contention** — purely changes which bytes are at which file offsets
|
| 101 |
+
- The only risk: co-occurrence patterns change with different prompts.
|
| 102 |
+
If the clustering is prompt-dependent, it might help some prompts and hurt others.
|
| 103 |
+
- Mitigation: use a diverse set of prompts for profiling
|
| 104 |
+
|
| 105 |
+
### Expected impact
|
| 106 |
+
- Best case: 30-50% expert_io reduction for cold reads (scattered → near-sequential)
|
| 107 |
+
- Worst case: neutral (if co-occurrence is too flat/prompt-dependent)
|
| 108 |
+
- Pipeline impact: pure I/O improvement, no GPU/CPU interaction
|
| 109 |
+
|
| 110 |
+
---
|
| 111 |
+
|
| 112 |
+
## Experiment 4: LZ4 DRAM expert cache
|
| 113 |
+
|
| 114 |
+
### Isolated test
|
| 115 |
+
- Allocate 4GB of malloc'd memory
|
| 116 |
+
- After each expert pread, LZ4-compress and store in the cache (hash by layer+expert_id)
|
| 117 |
+
- On subsequent reads: check cache first. Hit = LZ4 decompress from DRAM.
|
| 118 |
+
- Measure: cache hit rate over 200 tokens, avg expert read time (hit vs miss)
|
| 119 |
+
|
| 120 |
+
### What this does
|
| 121 |
+
- Creates a second-level cache between the OS page cache and SSD
|
| 122 |
+
- Stores experts in compressed form → 4GB holds ~730 experts (vs ~570 raw)
|
| 123 |
+
- Cache hit: decompress at 41 GB/s = 0.17ms per expert
|
| 124 |
+
- Cache miss: pread from SSD/page cache = 0.3-1.5ms per expert
|
| 125 |
+
- The cache is in USERSPACE memory — doesn't interfere with OS page cache
|
| 126 |
+
|
| 127 |
+
### Pipeline contention analysis
|
| 128 |
+
- Cache lookup: hash table check (~0.001ms) — negligible
|
| 129 |
+
- Cache hit: LZ4 decompress runs on the I/O worker thread (CPU)
|
| 130 |
+
- During decompress, the CPU is busy for 0.17ms
|
| 131 |
+
- This overlaps with other threads' preads (parallel, no contention)
|
| 132 |
+
- Cache miss: normal pread path (no change)
|
| 133 |
+
- **Memory contention**: 4GB of malloc'd DRAM reduces available page cache by 4GB
|
| 134 |
+
(OS has 4GB less to work with). Current page cache: ~35GB → ~31GB.
|
| 135 |
+
31GB / 7MB = ~4430 expert slots (vs current 5000). FEWER raw experts cached.
|
| 136 |
+
BUT the userspace cache adds 730 compressed experts on top.
|
| 137 |
+
Total accessible experts: 4430 (page cache) + 730 (LZ4 cache) = 5160.
|
| 138 |
+
That's only 3% more than current 5000. Barely worth it.
|
| 139 |
+
- **If we increase to 8GB cache**: 27GB page cache (3857 slots) + 1455 LZ4 = 5312. Still marginal.
|
| 140 |
+
- **The math doesn't work** unless the LZ4 cache has MUCH higher hit rate than the
|
| 141 |
+
page cache (e.g., by using a smarter eviction policy than LRU).
|
| 142 |
+
|
| 143 |
+
### Expected impact
|
| 144 |
+
- Best case: 5-10% expert_io reduction (LZ4 cache hits for hottest experts)
|
| 145 |
+
- Worst case: negative (reduced page cache hurts more than LZ4 cache helps)
|
| 146 |
+
- Pipeline impact: CPU time for decompress, but overlaps with parallel preads
|
| 147 |
+
- **VERDICT: probably not worth implementing given the math above**
|
| 148 |
+
|
| 149 |
+
---
|
| 150 |
+
|
| 151 |
+
## Experiment 5: aio_read batching
|
| 152 |
+
|
| 153 |
+
### Isolated test
|
| 154 |
+
- (A) 4 × pread on 4 pthreads (current)
|
| 155 |
+
- (B) 4 × aio_read, then aio_suspend to wait for all
|
| 156 |
+
- Both F_NOCACHE. Both warm cache. 50 iterations.
|
| 157 |
+
|
| 158 |
+
### What aio_read does differently
|
| 159 |
+
- Submits I/O requests to the kernel without blocking the calling thread
|
| 160 |
+
- The kernel sees all 4 requests at once and can batch NVMe commands
|
| 161 |
+
- aio_suspend blocks until all 4 complete (single wait vs 4 thread joins)
|
| 162 |
+
- Eliminates per-thread overhead (no pthread_create/join or dispatch_group)
|
| 163 |
+
|
| 164 |
+
### Pipeline contention analysis
|
| 165 |
+
- aio_read uses kernel-level async I/O (not userspace threads)
|
| 166 |
+
- The kernel's I/O scheduler has full visibility into all pending reads
|
| 167 |
+
- **No new contention** — same SSD path, potentially better NVMe scheduling
|
| 168 |
+
- Risk: macOS aio implementation might be less optimized than GCD
|
| 169 |
+
(Apple generally prefers dispatch_io over POSIX aio)
|
| 170 |
+
- The completion notification (SIGEV_THREAD or SIGEV_SIGNAL) has latency
|
| 171 |
+
|
| 172 |
+
### Expected impact
|
| 173 |
+
- Best case: 5-15% expert_io reduction (better NVMe command batching)
|
| 174 |
+
- Worst case: neutral or regression (aio overhead on macOS)
|
| 175 |
+
- Pipeline impact: drop-in replacement for pread
|
| 176 |
+
|
| 177 |
+
---
|
| 178 |
+
|
| 179 |
+
## Execution Priority
|
| 180 |
+
|
| 181 |
+
| # | Experiment | Expected Impact | Effort | Risk | Priority |
|
| 182 |
+
|---|-----------|----------------|--------|------|----------|
|
| 183 |
+
| 3 | Expert clustering | 30-50% cold I/O | Medium | Low | **1st** |
|
| 184 |
+
| 1 | dispatch_io | 10-20% I/O | Low | Low | **2nd** |
|
| 185 |
+
| 2 | GPU private compression | 15-30% GPU | Medium | Medium | **3rd** |
|
| 186 |
+
| 5 | aio_read | 5-15% I/O | Low | Low | **4th** |
|
| 187 |
+
| 4 | LZ4 DRAM cache | 5-10% I/O | High | High | Skip |
|
| 188 |
+
|
| 189 |
+
## Key Prediction: Compound Effects
|
| 190 |
+
|
| 191 |
+
The reason to test each in isolation FIRST: on unified memory, improvements that look
|
| 192 |
+
good alone can cancel each other (like F_RDADVISE's expert_io -31% + cmd2_wait +73% = net 0%).
|
| 193 |
+
|
| 194 |
+
After isolated tests, the compound analysis:
|
| 195 |
+
- Experiment 1 + 3 (dispatch_io + clustering): better I/O scheduling + fewer scattered reads.
|
| 196 |
+
These should compound because clustering reduces scatter and dispatch_io optimizes remaining scatter.
|
| 197 |
+
- Experiment 2 is GPU-only, orthogonal to I/O experiments. Should compound cleanly.
|
| 198 |
+
- Experiment 5 is alternative to experiment 1 — test both, pick winner.
|
docs/quantization-guide.md
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Quantization Guide: DWQ, JANG, and the GPTQ Pipeline
|
| 2 |
+
|
| 3 |
+
This document covers the Flash-MoE quantization pipeline, including the DWQ (Data-aware Weight Quantization) and JANG (Jang Adaptive N-bit Grading) approaches, their relationship, and the full GPTQ-based implementation.
|
| 4 |
+
|
| 5 |
+
## DWQ vs JANG
|
| 6 |
+
|
| 7 |
+
DWQ and JANG solve different problems and are **complementary, not competing**. JANG decides HOW MANY bits each tensor gets. DWQ decides HOW to use those bits optimally.
|
| 8 |
+
|
| 9 |
+
### What is DWQ (Data-aware Weight Quantization)?
|
| 10 |
+
|
| 11 |
+
DWQ is GPTQ-style calibration that uses real input statistics (the Hessian proxy H = X^T @ X) to guide error compensation during quantization. Instead of independently rounding each weight to its nearest quantization level (RTN -- Round To Nearest), DWQ processes weights column-by-column and compensates downstream columns for each rounding error. The compensation is weighted by input importance: weights that participate in high-variance input dimensions receive more aggressive error correction.
|
| 12 |
+
|
| 13 |
+
The result: same number of bits, but the quantized matrix produces outputs much closer to the original.
|
| 14 |
+
|
| 15 |
+
### What is JANG (Jang Adaptive N-bit Grading)?
|
| 16 |
+
|
| 17 |
+
JANG assigns different bit widths to different tensors or layer types based on their sensitivity to quantization error. The core insight is that not all weights are equally important:
|
| 18 |
+
|
| 19 |
+
- **Router weights** (expert gating): extremely sensitive to quantization -- small errors flip expert selection entirely. Assigned **8-bit**.
|
| 20 |
+
- **Attention projections** (Q, K, V, O): moderate sensitivity. Assigned **4-bit**.
|
| 21 |
+
- **Hot experts** (top ~25% by activation frequency): handle ~80% of tokens. Assigned **4-bit**.
|
| 22 |
+
- **Cold experts** (remaining ~75%): infrequently activated. Assigned **2-bit**.
|
| 23 |
+
|
| 24 |
+
### Comparison
|
| 25 |
+
|
| 26 |
+
| | DWQ (GPTQ-style) | JANG (Adaptive N-bit) |
|
| 27 |
+
|---|---|---|
|
| 28 |
+
| **What it decides** | How to quantize (optimal scale/bias values) | How many bits per tensor |
|
| 29 |
+
| **Error reduction type** | Mathematical -- minimizes output reconstruction error | Structural -- allocates bits where they matter most |
|
| 30 |
+
| **Our 2-bit problem** | Fixes broken JSON at 2-bit by compensating rounding errors | Avoids the problem by keeping sensitive tensors at higher bit widths |
|
| 31 |
+
| **Model size impact** | None -- same bit width, better values | Large -- mixed precision reduces total size |
|
| 32 |
+
| **Offline cost** | Moderate -- needs calibration data + Hessian computation | Low -- needs expert frequency profiling |
|
| 33 |
+
|
| 34 |
+
### Why Neither Alone Is Sufficient
|
| 35 |
+
|
| 36 |
+
- **JANG 2-bit without DWQ**: Cold experts at 2-bit with naive RTN quantization still produce broken JSON. The `\name\` instead of `"name"` problem persists because RTN 2-bit has too much uncorrelated rounding error.
|
| 37 |
+
- **DWQ 4-bit without JANG**: Every expert stays at 4-bit. The model is 209GB. No size reduction.
|
| 38 |
+
|
| 39 |
+
### The Ideal Pipeline
|
| 40 |
+
|
| 41 |
+
```
|
| 42 |
+
JANG first (architecture decision) --> DWQ/GPTQ second (math optimization)
|
| 43 |
+
"Which tensors get 2-bit?" "Make those 2-bit values as good as possible"
|
| 44 |
+
```
|
| 45 |
+
|
| 46 |
+
JANG is the structural decision: assign bit budgets based on sensitivity. DWQ is the mathematical optimization: within each bit budget, find the quantized values that minimize output error. Together, they produce a model that is both smaller (JANG) and higher quality (DWQ).
|
| 47 |
+
|
| 48 |
+
## Quantization Formats
|
| 49 |
+
|
| 50 |
+
### 4-bit Format
|
| 51 |
+
|
| 52 |
+
- 8 values packed per `uint32`, LSB-first (bits 0-3 = value 0, bits 4-7 = value 1, etc.)
|
| 53 |
+
- Scale and bias per group of 64 values, stored as `bf16`
|
| 54 |
+
- Dequantization: `value = uint4_nibble * scale + bias`
|
| 55 |
+
|
| 56 |
+
### 2-bit Format
|
| 57 |
+
|
| 58 |
+
- 16 values packed per `uint32`, LSB-first (bits 0-1 = value 0, bits 2-3 = value 1, etc.)
|
| 59 |
+
- Scale and bias per group of 64 values, stored as `bf16`
|
| 60 |
+
- Dequantization: `value = uint2_pair * scale + bias`
|
| 61 |
+
|
| 62 |
+
### bf16 Conversion
|
| 63 |
+
|
| 64 |
+
Scale and bias are stored as `bf16` (bfloat16). Conversion from `float32` is a simple truncation: `bf16 = float32_bits >> 16`. No rounding is applied.
|
| 65 |
+
|
| 66 |
+
### RTN vs GPTQ: Same Format, Different Values
|
| 67 |
+
|
| 68 |
+
The on-disk format is identical for RTN and GPTQ quantized weights. Both produce the same packed uint4/uint2 values with bf16 scale and bias per group. The Metal dequant formula is the same. The only difference is in the scale/bias values themselves -- GPTQ produces values that minimize output reconstruction error rather than per-weight rounding error.
|
| 69 |
+
|
| 70 |
+
## Our Implementation: 4-Phase Pipeline
|
| 71 |
+
|
| 72 |
+
### Phase 0: MSE-Optimal Clipping
|
| 73 |
+
|
| 74 |
+
Before any GPTQ calibration, we apply optimal clipping to the RTN baseline. For each group of 64 values, a grid search over 20 clipping ratios finds the ratio that minimizes mean squared error between the original float32 values and their quantized reconstruction. This alone achieves 15-30% RMSE reduction over naive min/max scaling.
|
| 75 |
+
|
| 76 |
+
Implemented in `repack_experts_2bit.py`.
|
| 77 |
+
|
| 78 |
+
### Phase 1: Calibration Collection
|
| 79 |
+
|
| 80 |
+
Collect real input activations for each expert to build the Hessian proxy.
|
| 81 |
+
|
| 82 |
+
1. Run inference with `--collect-activations` flag. This dumps expert input vectors (the x fed into each expert MLP) to disk.
|
| 83 |
+
2. `build_hessian.py` accumulates H = X^T @ X per expert in an online fashion (no need to store all activations in memory).
|
| 84 |
+
3. `calibrate.sh` orchestrates the collection across diverse prompts. Minimum 16K tokens recommended for stable Hessian estimates.
|
| 85 |
+
|
| 86 |
+
### Phase 2: GPTQ Requantization
|
| 87 |
+
|
| 88 |
+
`gptq_requantize.py` applies the blocked GPTQ algorithm to requantize experts from 4-bit to 2-bit using the collected Hessians.
|
| 89 |
+
|
| 90 |
+
Key details:
|
| 91 |
+
- **Block size**: 128 columns. Processing in blocks amortizes the Cholesky factorization cost and improves numerical stability.
|
| 92 |
+
- **Automatic fallback**: Experts without calibration data (never activated during calibration) fall back to MSE-optimal clipping from Phase 0.
|
| 93 |
+
- **Safety check**: GPTQ output is only used if it beats RTN RMSE. If the Hessian is degenerate or the calibration data is insufficient, the RTN result is kept.
|
| 94 |
+
|
| 95 |
+
### Phase 3: Sensitivity Analysis
|
| 96 |
+
|
| 97 |
+
`sensitivity_analysis.py` computes a sensitivity score per expert:
|
| 98 |
+
|
| 99 |
+
```
|
| 100 |
+
sensitivity = frequency * quant_error * layer_weight
|
| 101 |
+
```
|
| 102 |
+
|
| 103 |
+
- `frequency`: how often the expert is activated (from profiling)
|
| 104 |
+
- `quant_error`: RMSE between float32 and quantized output
|
| 105 |
+
- `layer_weight`: earlier layers are weighted higher (errors compound through the network)
|
| 106 |
+
|
| 107 |
+
Experts are ranked by sensitivity. The most sensitive experts are assigned 4-bit until the target disk budget (in GB) is reached. The rest get GPTQ 2-bit. The output is a `hot_experts.json` manifest consumed by `repack_experts_tiered.py`.
|
| 108 |
+
|
| 109 |
+
The updated `repack_experts_tiered.py` accepts `--gptq-dir` (directory of GPTQ-requantized 2-bit experts) and `--hot-experts` (JSON manifest from sensitivity analysis) to build the final tiered model.
|
| 110 |
+
|
| 111 |
+
## GPTQ Algorithm (Technical Details)
|
| 112 |
+
|
| 113 |
+
The GPTQ algorithm minimizes the layer-wise reconstruction error:
|
| 114 |
+
|
| 115 |
+
```
|
| 116 |
+
argmin_Q || W @ X - Q @ X ||^2
|
| 117 |
+
```
|
| 118 |
+
|
| 119 |
+
where W is the original weight matrix, Q is the quantized matrix, and X is the calibration input. This is equivalent to minimizing the error weighted by the Hessian H = X^T @ X.
|
| 120 |
+
|
| 121 |
+
### Step-by-Step
|
| 122 |
+
|
| 123 |
+
1. **Collect calibration activations** X per expert (Phase 1).
|
| 124 |
+
2. **Compute Hessian proxy** H = X^T @ X. This is a d_in x d_in matrix capturing input correlations.
|
| 125 |
+
3. **Add damping**: H += lambda * mean(diag(H)) * I. Prevents numerical instability from near-zero eigenvalues.
|
| 126 |
+
4. **Cholesky factorize** H to get H_inv efficiently.
|
| 127 |
+
5. **Process columns left-to-right in blocks of 128**:
|
| 128 |
+
- For each column j in the block:
|
| 129 |
+
- Quantize w_j to the nearest quantization level: q_j = quantize(w_j)
|
| 130 |
+
- Compute the quantization error: e_j = w_j - q_j
|
| 131 |
+
- Compensate remaining columns: w_{j+1..n} += e_j * H_inv[j, j+1..n] / H_inv[j, j]
|
| 132 |
+
6. **Net effect**: Errors in important input dimensions (high H diagonal values) are compensated more aggressively. Weights connected to correlated inputs are adjusted together.
|
| 133 |
+
|
| 134 |
+
### Why GPTQ Works Better Than RTN
|
| 135 |
+
|
| 136 |
+
RTN (Round To Nearest) independently rounds each weight to the nearest quantization level. This is locally optimal per weight but globally suboptimal because it ignores correlations.
|
| 137 |
+
|
| 138 |
+
GPTQ accounts for weight-to-weight correlations through the Hessian. When weight w_j is rounded down, GPTQ adjusts w_{j+1..n} to compensate, with the adjustment magnitude guided by input statistics. Weights connected to high-variance inputs get larger corrections. The result: same number of bits, but the quantized matrix produces outputs much closer to the original when evaluated on real inputs.
|
| 139 |
+
|
| 140 |
+
At 4-bit, RTN is already quite good (the quantization grid is fine enough). At 2-bit, the grid is coarse (only 4 levels), and RTN's uncorrelated rounding errors accumulate catastrophically -- this is why RTN 2-bit breaks JSON output. GPTQ's error compensation keeps the accumulated error bounded.
|
| 141 |
+
|
| 142 |
+
## Full Pipeline Commands
|
| 143 |
+
|
| 144 |
+
```bash
|
| 145 |
+
# Step 1: Calibration -- collect expert input activations
|
| 146 |
+
cd metal_infer && ./calibrate.sh
|
| 147 |
+
|
| 148 |
+
# Step 2: GPTQ requantization -- apply blocked GPTQ to 2-bit experts
|
| 149 |
+
python gptq_requantize.py --hessian-dir calibration/ --parallel 8
|
| 150 |
+
|
| 151 |
+
# Step 3: Sensitivity analysis -- rank experts, decide 4-bit vs 2-bit
|
| 152 |
+
python sensitivity_analysis.py --packed-dir ../packed_experts/ --target-gb 150
|
| 153 |
+
|
| 154 |
+
# Step 4: Build tiered model -- assemble final mixed-precision expert pack
|
| 155 |
+
python ../repack_experts_tiered.py --hot-experts hot_experts.json --gptq-dir packed_experts_gptq_2bit/
|
| 156 |
+
```
|
| 157 |
+
|
| 158 |
+
### What Each Step Produces
|
| 159 |
+
|
| 160 |
+
| Step | Input | Output |
|
| 161 |
+
|------|-------|--------|
|
| 162 |
+
| calibrate.sh | Diverse prompts + inference engine | `calibration/` directory with per-expert Hessian matrices |
|
| 163 |
+
| gptq_requantize.py | 4-bit experts + Hessians | `packed_experts_gptq_2bit/` directory |
|
| 164 |
+
| sensitivity_analysis.py | Packed experts + frequency data | `hot_experts.json` manifest |
|
| 165 |
+
| repack_experts_tiered.py | 4-bit experts + GPTQ 2-bit experts + manifest | `packed_experts_tiered/` directory (production model) |
|
| 166 |
+
|
| 167 |
+
## Model Size Estimates
|
| 168 |
+
|
| 169 |
+
| Configuration | Size | Quality | Notes |
|
| 170 |
+
|---------------|------|---------|-------|
|
| 171 |
+
| All 4-bit | 209 GB | Excellent | Current production. Full tool calling. |
|
| 172 |
+
| All RTN 2-bit | 120 GB | Broken JSON | `\name\` instead of `"name"` in JSON output |
|
| 173 |
+
| All GPTQ 2-bit | 120 GB | Good (expected) | GPTQ error compensation fixes JSON |
|
| 174 |
+
| 20% hot 4-bit + 80% GPTQ 2-bit | 134 GB | Very good (expected) | Best size/quality tradeoff |
|
| 175 |
+
|
| 176 |
+
The 20/80 tiered configuration with GPTQ is the target for production deployment on storage-constrained devices (256GB iPhones, smaller SSDs).
|
| 177 |
+
|
| 178 |
+
## Files
|
| 179 |
+
|
| 180 |
+
| File | Location | Purpose |
|
| 181 |
+
|------|----------|---------|
|
| 182 |
+
| `gptq_requantize.py` | `metal_infer/` | Blocked GPTQ requantization with automatic RTN fallback |
|
| 183 |
+
| `build_hessian.py` | `metal_infer/` | Online Hessian accumulation (H = X^T @ X) per expert |
|
| 184 |
+
| `sensitivity_analysis.py` | `metal_infer/` | Expert sensitivity scoring and bit-width assignment |
|
| 185 |
+
| `calibrate.sh` | `metal_infer/` | Calibration runner (diverse prompts, 16K+ tokens) |
|
| 186 |
+
| `repack_experts_2bit.py` | `metal_infer/` | RTN 2-bit repacking with MSE-optimal clipping (Phase 0) |
|
| 187 |
+
| `repack_experts_tiered.py` | root | Tiered repacking with `--gptq-dir` and `--hot-experts` support |
|
docs/superpowers/plans/2026-03-20-runtime-model-config.md
ADDED
|
@@ -0,0 +1,847 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Runtime Model Config Implementation Plan
|
| 2 |
+
|
| 3 |
+
> **For agentic workers:** REQUIRED: Use superpowers:subagent-driven-development (if subagents available) or superpowers:executing-plans to implement this plan. Steps use checkbox (`- [ ]`) syntax for tracking.
|
| 4 |
+
|
| 5 |
+
**Goal:** Replace all ~54 model-specific `#define` constants in `infer.m` with a runtime `ModelConfig` struct populated from HuggingFace `config.json` + `tokenizer.json`, enabling model switching via `--model` flag without recompilation.
|
| 6 |
+
|
| 7 |
+
**Architecture:** A single `ModelConfig cfg` global struct is populated at startup by parsing JSON files using NSJSONSerialization. Expert byte offsets are computed from dimensions + quantization params. Static arrays sized by model constants become dynamically allocated after config loading.
|
| 8 |
+
|
| 9 |
+
**Tech Stack:** Objective-C (NSJSONSerialization), C (malloc/free), Metal (unchanged — already parameterized)
|
| 10 |
+
|
| 11 |
+
**Spec:** `docs/superpowers/specs/2026-03-20-runtime-model-config-design.md`
|
| 12 |
+
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
## File Structure
|
| 16 |
+
|
| 17 |
+
All changes are in a single file:
|
| 18 |
+
|
| 19 |
+
- **Modify:** `metal_infer/infer.m` — the entire inference engine (~7200 lines)
|
| 20 |
+
- Remove lines 68-140 (model `#define`s)
|
| 21 |
+
- Add `ModelConfig` struct + `load_model_config()` + `compute_expert_offsets()` + `alloc_tracking_arrays()` (~200 lines)
|
| 22 |
+
- Convert ~13 static/stack arrays to dynamic allocation
|
| 23 |
+
- Replace ~960 occurrences of `#define` names with `cfg.field` references
|
| 24 |
+
|
| 25 |
+
No other files need changes. `shaders.metal` is already parameterized. `chat.m` is a pure HTTP client with no model constants.
|
| 26 |
+
|
| 27 |
+
---
|
| 28 |
+
|
| 29 |
+
## Task 1: Add ModelConfig struct and config loader
|
| 30 |
+
|
| 31 |
+
**Files:**
|
| 32 |
+
- Modify: `metal_infer/infer.m:68-140`
|
| 33 |
+
|
| 34 |
+
This task replaces the `#define` block with the `ModelConfig` struct, adds the JSON parsing function, expert offset computation, and dynamic array allocation. The old `#define`s are kept temporarily as fallback validation.
|
| 35 |
+
|
| 36 |
+
- [ ] **Step 1.1: Add ModelConfig struct after includes (before old defines)**
|
| 37 |
+
|
| 38 |
+
Insert after line 66 (after `#include <compression.h>`), before the old defines block. The struct holds all model-specific parameters:
|
| 39 |
+
|
| 40 |
+
```c
|
| 41 |
+
// ============================================================================
|
| 42 |
+
// Runtime model configuration (populated from HuggingFace config.json)
|
| 43 |
+
// ============================================================================
|
| 44 |
+
|
| 45 |
+
typedef struct {
|
| 46 |
+
// Core architecture
|
| 47 |
+
int hidden_dim;
|
| 48 |
+
int num_layers;
|
| 49 |
+
int num_attn_heads;
|
| 50 |
+
int num_kv_heads;
|
| 51 |
+
int head_dim;
|
| 52 |
+
int vocab_size;
|
| 53 |
+
float rms_norm_eps;
|
| 54 |
+
|
| 55 |
+
// MoE
|
| 56 |
+
int num_experts;
|
| 57 |
+
int num_experts_per_tok;
|
| 58 |
+
int moe_intermediate;
|
| 59 |
+
int shared_intermediate;
|
| 60 |
+
int group_size;
|
| 61 |
+
int bits;
|
| 62 |
+
|
| 63 |
+
// Linear attention (GatedDeltaNet)
|
| 64 |
+
int linear_num_v_heads;
|
| 65 |
+
int linear_num_k_heads;
|
| 66 |
+
int linear_key_dim;
|
| 67 |
+
int linear_value_dim;
|
| 68 |
+
int conv_kernel_size;
|
| 69 |
+
|
| 70 |
+
// Full attention
|
| 71 |
+
float rope_theta;
|
| 72 |
+
float partial_rotary;
|
| 73 |
+
|
| 74 |
+
// Layer type map
|
| 75 |
+
int num_full_attn_layers;
|
| 76 |
+
int num_linear_layers;
|
| 77 |
+
bool *is_full_attn; // [num_layers]
|
| 78 |
+
int *full_attn_index; // [num_layers] — index into full-attn buffers, or -1
|
| 79 |
+
int *linear_index; // [num_layers] — index into linear-attn buffers, or -1
|
| 80 |
+
|
| 81 |
+
// Derived: expert byte offsets (4-bit)
|
| 82 |
+
size_t expert_size_4bit;
|
| 83 |
+
size_t gate_w_off_4, gate_s_off_4, gate_b_off_4;
|
| 84 |
+
size_t up_w_off_4, up_s_off_4, up_b_off_4;
|
| 85 |
+
size_t down_w_off_4, down_s_off_4, down_b_off_4;
|
| 86 |
+
|
| 87 |
+
// Derived: expert byte offsets (2-bit)
|
| 88 |
+
size_t expert_size_2bit;
|
| 89 |
+
size_t gate_w_off_2, gate_s_off_2, gate_b_off_2;
|
| 90 |
+
size_t up_w_off_2, up_s_off_2, up_b_off_2;
|
| 91 |
+
size_t down_w_off_2, down_s_off_2, down_b_off_2;
|
| 92 |
+
|
| 93 |
+
// Derived dimensions
|
| 94 |
+
int linear_total_key;
|
| 95 |
+
int linear_total_value;
|
| 96 |
+
int linear_conv_dim;
|
| 97 |
+
int rotary_dim;
|
| 98 |
+
|
| 99 |
+
// Special tokens
|
| 100 |
+
int eos_token_ids[8];
|
| 101 |
+
int num_eos_tokens;
|
| 102 |
+
int think_start_token;
|
| 103 |
+
int think_end_token;
|
| 104 |
+
|
| 105 |
+
// Context limits
|
| 106 |
+
int max_seq_len;
|
| 107 |
+
int gpu_kv_seq;
|
| 108 |
+
|
| 109 |
+
// Model path (resolved)
|
| 110 |
+
char model_path[1024];
|
| 111 |
+
} ModelConfig;
|
| 112 |
+
|
| 113 |
+
static ModelConfig cfg;
|
| 114 |
+
```
|
| 115 |
+
|
| 116 |
+
- [ ] **Step 1.2: Add compute_expert_offsets()**
|
| 117 |
+
|
| 118 |
+
Place right after the struct definition:
|
| 119 |
+
|
| 120 |
+
```c
|
| 121 |
+
static void compute_expert_offsets(ModelConfig *c) {
|
| 122 |
+
int mid = c->moe_intermediate;
|
| 123 |
+
int hid = c->hidden_dim;
|
| 124 |
+
int gs = c->group_size;
|
| 125 |
+
|
| 126 |
+
for (int b = 4; b >= 2; b -= 2) {
|
| 127 |
+
int vals_per_u32 = 32 / b;
|
| 128 |
+
// gate_proj [mid, hid]
|
| 129 |
+
size_t gw = (size_t)mid * ((hid + vals_per_u32 - 1) / vals_per_u32) * 4;
|
| 130 |
+
size_t gs_sz = (size_t)mid * ((hid + gs - 1) / gs) * 2;
|
| 131 |
+
size_t gb = gs_sz;
|
| 132 |
+
// up_proj [mid, hid] — same shape
|
| 133 |
+
size_t uw = gw, us = gs_sz, ub = gb;
|
| 134 |
+
// down_proj [hid, mid]
|
| 135 |
+
size_t dw = (size_t)hid * ((mid + vals_per_u32 - 1) / vals_per_u32) * 4;
|
| 136 |
+
size_t ds = (size_t)hid * ((mid + gs - 1) / gs) * 2;
|
| 137 |
+
size_t db = ds;
|
| 138 |
+
|
| 139 |
+
size_t off = 0;
|
| 140 |
+
if (b == 4) {
|
| 141 |
+
c->gate_w_off_4 = off; off += gw;
|
| 142 |
+
c->gate_s_off_4 = off; off += gs_sz;
|
| 143 |
+
c->gate_b_off_4 = off; off += gb;
|
| 144 |
+
c->up_w_off_4 = off; off += uw;
|
| 145 |
+
c->up_s_off_4 = off; off += us;
|
| 146 |
+
c->up_b_off_4 = off; off += ub;
|
| 147 |
+
c->down_w_off_4 = off; off += dw;
|
| 148 |
+
c->down_s_off_4 = off; off += ds;
|
| 149 |
+
c->down_b_off_4 = off; off += db;
|
| 150 |
+
c->expert_size_4bit = off;
|
| 151 |
+
} else {
|
| 152 |
+
c->gate_w_off_2 = off; off += gw;
|
| 153 |
+
c->gate_s_off_2 = off; off += gs_sz;
|
| 154 |
+
c->gate_b_off_2 = off; off += gb;
|
| 155 |
+
c->up_w_off_2 = off; off += uw;
|
| 156 |
+
c->up_s_off_2 = off; off += us;
|
| 157 |
+
c->up_b_off_2 = off; off += ub;
|
| 158 |
+
c->down_w_off_2 = off; off += dw;
|
| 159 |
+
c->down_s_off_2 = off; off += ds;
|
| 160 |
+
c->down_b_off_2 = off; off += db;
|
| 161 |
+
c->expert_size_2bit = off;
|
| 162 |
+
}
|
| 163 |
+
}
|
| 164 |
+
}
|
| 165 |
+
```
|
| 166 |
+
|
| 167 |
+
- [ ] **Step 1.3: Add load_model_config() — the JSON parser**
|
| 168 |
+
|
| 169 |
+
This function reads `config.json` and `tokenizer.json` using NSJSONSerialization:
|
| 170 |
+
|
| 171 |
+
```c
|
| 172 |
+
static void load_model_config(const char *model_dir) {
|
| 173 |
+
memset(&cfg, 0, sizeof(cfg));
|
| 174 |
+
cfg.think_start_token = -1;
|
| 175 |
+
cfg.think_end_token = -1;
|
| 176 |
+
cfg.gpu_kv_seq = 8192;
|
| 177 |
+
|
| 178 |
+
// Resolve HF snapshot directory
|
| 179 |
+
NSString *base = [NSString stringWithUTF8String:model_dir];
|
| 180 |
+
NSString *configPath = [base stringByAppendingPathComponent:@"config.json"];
|
| 181 |
+
NSFileManager *fm = [NSFileManager defaultManager];
|
| 182 |
+
|
| 183 |
+
if (![fm fileExistsAtPath:configPath]) {
|
| 184 |
+
NSString *snapDir = [base stringByAppendingPathComponent:@"snapshots"];
|
| 185 |
+
if ([fm fileExistsAtPath:snapDir]) {
|
| 186 |
+
NSArray *snaps = [[fm contentsOfDirectoryAtPath:snapDir error:nil]
|
| 187 |
+
sortedArrayUsingSelector:@selector(compare:)];
|
| 188 |
+
for (NSString *snap in snaps) {
|
| 189 |
+
NSString *candidate = [[snapDir stringByAppendingPathComponent:snap]
|
| 190 |
+
stringByAppendingPathComponent:@"config.json"];
|
| 191 |
+
if ([fm fileExistsAtPath:candidate]) {
|
| 192 |
+
base = [snapDir stringByAppendingPathComponent:snap];
|
| 193 |
+
configPath = candidate;
|
| 194 |
+
break;
|
| 195 |
+
}
|
| 196 |
+
}
|
| 197 |
+
}
|
| 198 |
+
}
|
| 199 |
+
|
| 200 |
+
if (![fm fileExistsAtPath:configPath]) {
|
| 201 |
+
fprintf(stderr, "FATAL: config.json not found in %s\n", model_dir);
|
| 202 |
+
exit(1);
|
| 203 |
+
}
|
| 204 |
+
|
| 205 |
+
strlcpy(cfg.model_path, [base UTF8String], sizeof(cfg.model_path));
|
| 206 |
+
|
| 207 |
+
// Parse config.json
|
| 208 |
+
NSData *data = [NSData dataWithContentsOfFile:configPath];
|
| 209 |
+
NSDictionary *root = [NSJSONSerialization JSONObjectWithData:data options:0 error:nil];
|
| 210 |
+
NSDictionary *tc = root[@"text_config"];
|
| 211 |
+
if (!tc) { fprintf(stderr, "FATAL: config.json missing text_config\n"); exit(1); }
|
| 212 |
+
|
| 213 |
+
cfg.hidden_dim = [tc[@"hidden_size"] intValue];
|
| 214 |
+
cfg.num_layers = [tc[@"num_hidden_layers"] intValue];
|
| 215 |
+
cfg.num_attn_heads = [tc[@"num_attention_heads"] intValue];
|
| 216 |
+
cfg.num_kv_heads = [tc[@"num_key_value_heads"] intValue];
|
| 217 |
+
cfg.head_dim = [tc[@"head_dim"] intValue];
|
| 218 |
+
cfg.vocab_size = [tc[@"vocab_size"] intValue];
|
| 219 |
+
cfg.rms_norm_eps = [tc[@"rms_norm_eps"] floatValue];
|
| 220 |
+
cfg.num_experts = [tc[@"num_experts"] intValue];
|
| 221 |
+
cfg.num_experts_per_tok = [tc[@"num_experts_per_tok"] intValue];
|
| 222 |
+
cfg.moe_intermediate = [tc[@"moe_intermediate_size"] intValue];
|
| 223 |
+
cfg.shared_intermediate = [tc[@"shared_expert_intermediate_size"] intValue];
|
| 224 |
+
cfg.linear_num_v_heads = [tc[@"linear_num_value_heads"] intValue];
|
| 225 |
+
cfg.linear_num_k_heads = [tc[@"linear_num_key_heads"] intValue];
|
| 226 |
+
cfg.linear_key_dim = [tc[@"linear_key_head_dim"] intValue];
|
| 227 |
+
cfg.linear_value_dim = [tc[@"linear_value_head_dim"] intValue];
|
| 228 |
+
cfg.conv_kernel_size = tc[@"linear_conv_kernel_dim"] ? [tc[@"linear_conv_kernel_dim"] intValue] : 4;
|
| 229 |
+
cfg.max_seq_len = [tc[@"max_position_embeddings"] intValue];
|
| 230 |
+
|
| 231 |
+
// Quantization
|
| 232 |
+
NSDictionary *qc = root[@"quantization_config"] ?: root[@"quantization"];
|
| 233 |
+
if (qc) {
|
| 234 |
+
cfg.group_size = [qc[@"group_size"] intValue];
|
| 235 |
+
cfg.bits = [qc[@"bits"] intValue];
|
| 236 |
+
} else {
|
| 237 |
+
cfg.group_size = 64;
|
| 238 |
+
cfg.bits = 4;
|
| 239 |
+
fprintf(stderr, "[config] WARNING: no quantization_config, defaulting to 4-bit group_size=64\n");
|
| 240 |
+
}
|
| 241 |
+
|
| 242 |
+
// RoPE parameters
|
| 243 |
+
NSDictionary *rope = tc[@"rope_parameters"];
|
| 244 |
+
if (rope) {
|
| 245 |
+
cfg.rope_theta = [rope[@"rope_theta"] floatValue];
|
| 246 |
+
cfg.partial_rotary = [rope[@"partial_rotary_factor"] floatValue];
|
| 247 |
+
} else {
|
| 248 |
+
cfg.rope_theta = 10000000.0f;
|
| 249 |
+
cfg.partial_rotary = 0.25f;
|
| 250 |
+
}
|
| 251 |
+
|
| 252 |
+
// Layer types
|
| 253 |
+
NSArray *layerTypes = tc[@"layer_types"];
|
| 254 |
+
cfg.is_full_attn = calloc(cfg.num_layers, sizeof(bool));
|
| 255 |
+
cfg.full_attn_index = malloc(cfg.num_layers * sizeof(int));
|
| 256 |
+
cfg.linear_index = malloc(cfg.num_layers * sizeof(int));
|
| 257 |
+
|
| 258 |
+
int full_count = 0, linear_count = 0;
|
| 259 |
+
if (layerTypes && [layerTypes count] == (NSUInteger)cfg.num_layers) {
|
| 260 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 261 |
+
cfg.is_full_attn[i] = [layerTypes[i] isEqualToString:@"full_attention"];
|
| 262 |
+
}
|
| 263 |
+
} else {
|
| 264 |
+
// Fallback: use full_attn_interval pattern
|
| 265 |
+
int interval = tc[@"full_attention_interval"] ? [tc[@"full_attention_interval"] intValue] : 4;
|
| 266 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 267 |
+
cfg.is_full_attn[i] = ((i + 1) % interval == 0);
|
| 268 |
+
}
|
| 269 |
+
fprintf(stderr, "[config] Using full_attn_interval=%d (no explicit layer_types)\n", interval);
|
| 270 |
+
}
|
| 271 |
+
|
| 272 |
+
for (int i = 0; i < cfg.num_layers; i++) {
|
| 273 |
+
if (cfg.is_full_attn[i]) {
|
| 274 |
+
cfg.full_attn_index[i] = full_count++;
|
| 275 |
+
cfg.linear_index[i] = -1;
|
| 276 |
+
} else {
|
| 277 |
+
cfg.linear_index[i] = linear_count++;
|
| 278 |
+
cfg.full_attn_index[i] = -1;
|
| 279 |
+
}
|
| 280 |
+
}
|
| 281 |
+
cfg.num_full_attn_layers = full_count;
|
| 282 |
+
cfg.num_linear_layers = linear_count;
|
| 283 |
+
|
| 284 |
+
// EOS tokens (can be int or array in config.json)
|
| 285 |
+
id eosVal = root[@"eos_token_id"];
|
| 286 |
+
if ([eosVal isKindOfClass:[NSArray class]]) {
|
| 287 |
+
NSArray *arr = (NSArray *)eosVal;
|
| 288 |
+
cfg.num_eos_tokens = (int)[arr count];
|
| 289 |
+
if (cfg.num_eos_tokens > 8) cfg.num_eos_tokens = 8;
|
| 290 |
+
for (int i = 0; i < cfg.num_eos_tokens; i++)
|
| 291 |
+
cfg.eos_token_ids[i] = [arr[i] intValue];
|
| 292 |
+
} else {
|
| 293 |
+
cfg.num_eos_tokens = 1;
|
| 294 |
+
cfg.eos_token_ids[0] = [eosVal intValue];
|
| 295 |
+
}
|
| 296 |
+
|
| 297 |
+
// Think tokens from tokenizer.json added_tokens
|
| 298 |
+
NSString *tokPath = [base stringByAppendingPathComponent:@"tokenizer.json"];
|
| 299 |
+
if ([fm fileExistsAtPath:tokPath]) {
|
| 300 |
+
NSData *tokData = [NSData dataWithContentsOfFile:tokPath];
|
| 301 |
+
NSDictionary *tokRoot = [NSJSONSerialization JSONObjectWithData:tokData options:0 error:nil];
|
| 302 |
+
NSArray *addedTokens = tokRoot[@"added_tokens"];
|
| 303 |
+
if (addedTokens) {
|
| 304 |
+
for (NSDictionary *tok in addedTokens) {
|
| 305 |
+
NSString *content = tok[@"content"];
|
| 306 |
+
int tid = [tok[@"id"] intValue];
|
| 307 |
+
if ([content isEqualToString:@"<think>"]) cfg.think_start_token = tid;
|
| 308 |
+
else if ([content isEqualToString:@"</think>"]) cfg.think_end_token = tid;
|
| 309 |
+
}
|
| 310 |
+
}
|
| 311 |
+
} else {
|
| 312 |
+
fprintf(stderr, "[config] WARNING: tokenizer.json not found, think tokens disabled\n");
|
| 313 |
+
}
|
| 314 |
+
|
| 315 |
+
// Derived dimensions
|
| 316 |
+
cfg.linear_total_key = cfg.linear_num_k_heads * cfg.linear_key_dim;
|
| 317 |
+
cfg.linear_total_value = cfg.linear_num_v_heads * cfg.linear_value_dim;
|
| 318 |
+
cfg.linear_conv_dim = cfg.linear_total_key * 2 + cfg.linear_total_value;
|
| 319 |
+
cfg.rotary_dim = (int)(cfg.head_dim * cfg.partial_rotary);
|
| 320 |
+
|
| 321 |
+
// Expert byte offsets
|
| 322 |
+
compute_expert_offsets(&cfg);
|
| 323 |
+
|
| 324 |
+
// Summary
|
| 325 |
+
fprintf(stderr, "[config] %d layers (%d linear + %d full), hidden=%d, heads=%d, kv_heads=%d, head_dim=%d\n",
|
| 326 |
+
cfg.num_layers, cfg.num_linear_layers, cfg.num_full_attn_layers,
|
| 327 |
+
cfg.hidden_dim, cfg.num_attn_heads, cfg.num_kv_heads, cfg.head_dim);
|
| 328 |
+
fprintf(stderr, "[config] %d experts (K=%d), moe_intermediate=%d, shared=%d\n",
|
| 329 |
+
cfg.num_experts, cfg.num_experts_per_tok, cfg.moe_intermediate, cfg.shared_intermediate);
|
| 330 |
+
fprintf(stderr, "[config] %d-bit quantization, group_size=%d, expert_size=%zu bytes\n",
|
| 331 |
+
cfg.bits, cfg.group_size, cfg.expert_size_4bit);
|
| 332 |
+
fprintf(stderr, "[config] EOS tokens: [");
|
| 333 |
+
for (int i = 0; i < cfg.num_eos_tokens; i++)
|
| 334 |
+
fprintf(stderr, "%s%d", i ? ", " : "", cfg.eos_token_ids[i]);
|
| 335 |
+
fprintf(stderr, "], think: %d/%d\n", cfg.think_start_token, cfg.think_end_token);
|
| 336 |
+
}
|
| 337 |
+
```
|
| 338 |
+
|
| 339 |
+
- [ ] **Step 1.4: Add alloc_tracking_arrays()**
|
| 340 |
+
|
| 341 |
+
This replaces the static arrays that were sized by `#define` constants:
|
| 342 |
+
|
| 343 |
+
```c
|
| 344 |
+
// Dynamic tracking arrays (allocated after config is loaded)
|
| 345 |
+
static int *g_expert_freq = NULL;
|
| 346 |
+
static uint8_t *g_expert_seen = NULL;
|
| 347 |
+
static LZ4IndexEntry **g_lz4_index = NULL;
|
| 348 |
+
static uint8_t *g_cache_seen = NULL;
|
| 349 |
+
static uint64_t *g_cache_last_touch_token = NULL;
|
| 350 |
+
static uint64_t *g_cache_last_evict_token = NULL;
|
| 351 |
+
static int *g_pred_experts = NULL;
|
| 352 |
+
static int *g_pred_count = NULL;
|
| 353 |
+
|
| 354 |
+
// Helper macros for flattened 2D access
|
| 355 |
+
#define FREQ(l, e) g_expert_freq[(l) * cfg.num_experts + (e)]
|
| 356 |
+
#define EXPERT_SEEN_BYTE(l, e) g_expert_seen[(l) * ((cfg.num_experts + 7) / 8) + ((e) >> 3)]
|
| 357 |
+
#define CACHE_SEEN(l, e) g_cache_seen[(l) * cfg.num_experts + (e)]
|
| 358 |
+
#define CACHE_TOUCH(l, e) g_cache_last_touch_token[(l) * cfg.num_experts + (e)]
|
| 359 |
+
#define CACHE_EVICT(l, e) g_cache_last_evict_token[(l) * cfg.num_experts + (e)]
|
| 360 |
+
#define PRED_EXPERT(l, k) g_pred_experts[(l) * MAX_K + (k)]
|
| 361 |
+
#define PRED_COUNT(l) g_pred_count[(l)]
|
| 362 |
+
|
| 363 |
+
static void alloc_tracking_arrays(void) {
|
| 364 |
+
int nl = cfg.num_layers;
|
| 365 |
+
int ne = cfg.num_experts;
|
| 366 |
+
int seen_bytes_per_layer = (ne + 7) / 8;
|
| 367 |
+
|
| 368 |
+
g_expert_freq = calloc(nl * ne, sizeof(int));
|
| 369 |
+
g_expert_seen = calloc(nl * seen_bytes_per_layer, sizeof(uint8_t));
|
| 370 |
+
g_lz4_index = calloc(nl, sizeof(LZ4IndexEntry *));
|
| 371 |
+
g_cache_seen = calloc(nl * ne, sizeof(uint8_t));
|
| 372 |
+
g_cache_last_touch_token = calloc(nl * ne, sizeof(uint64_t));
|
| 373 |
+
g_cache_last_evict_token = calloc(nl * ne, sizeof(uint64_t));
|
| 374 |
+
g_pred_experts = calloc(nl * MAX_K, sizeof(int));
|
| 375 |
+
g_pred_count = calloc(nl, sizeof(int));
|
| 376 |
+
}
|
| 377 |
+
```
|
| 378 |
+
|
| 379 |
+
- [ ] **Step 1.5: Build and verify it compiles**
|
| 380 |
+
|
| 381 |
+
At this point the old `#define`s still exist (they'll be removed in Task 2). The new code coexists. Build to verify no syntax errors:
|
| 382 |
+
|
| 383 |
+
```bash
|
| 384 |
+
cd metal_infer && make clean && make
|
| 385 |
+
```
|
| 386 |
+
|
| 387 |
+
Expected: compiles successfully (old defines still in use, new code not yet called).
|
| 388 |
+
|
| 389 |
+
- [ ] **Step 1.6: Commit**
|
| 390 |
+
|
| 391 |
+
```bash
|
| 392 |
+
git add metal_infer/infer.m
|
| 393 |
+
git commit -m "feat: add ModelConfig struct and config loader (not yet wired up)"
|
| 394 |
+
```
|
| 395 |
+
|
| 396 |
+
---
|
| 397 |
+
|
| 398 |
+
## Task 2: Wire up config loading in main() and remove old defines
|
| 399 |
+
|
| 400 |
+
**Files:**
|
| 401 |
+
- Modify: `metal_infer/infer.m` — main() function (line ~6581) and defines block (lines 68-140)
|
| 402 |
+
|
| 403 |
+
- [ ] **Step 2.1: Call load_model_config() early in main()**
|
| 404 |
+
|
| 405 |
+
In `main()`, right after CLI arg parsing (after the switch block, around line 6654), add:
|
| 406 |
+
|
| 407 |
+
```c
|
| 408 |
+
// ---- Load model configuration from HF config.json ----
|
| 409 |
+
load_model_config(model_path);
|
| 410 |
+
alloc_tracking_arrays();
|
| 411 |
+
```
|
| 412 |
+
|
| 413 |
+
This must come BEFORE `metal_setup()` (line 6689) since Metal buffer allocations depend on cfg values.
|
| 414 |
+
|
| 415 |
+
- [ ] **Step 2.2: Remove MODEL_PATH_DEFAULT, use cfg.model_path**
|
| 416 |
+
|
| 417 |
+
Change line 6583 from:
|
| 418 |
+
```c
|
| 419 |
+
const char *model_path = MODEL_PATH_DEFAULT;
|
| 420 |
+
```
|
| 421 |
+
to:
|
| 422 |
+
```c
|
| 423 |
+
const char *model_path = NULL;
|
| 424 |
+
```
|
| 425 |
+
|
| 426 |
+
And add a default path fallback in the config loader if model_path is NULL:
|
| 427 |
+
```c
|
| 428 |
+
if (!model_dir || !model_dir[0]) {
|
| 429 |
+
// Try common HF cache locations
|
| 430 |
+
const char *home = getenv("HOME");
|
| 431 |
+
char probe[1024];
|
| 432 |
+
snprintf(probe, sizeof(probe), "%s/.cache/huggingface/hub", home);
|
| 433 |
+
// ... or just require --model
|
| 434 |
+
fprintf(stderr, "FATAL: --model path required\n");
|
| 435 |
+
exit(1);
|
| 436 |
+
}
|
| 437 |
+
```
|
| 438 |
+
|
| 439 |
+
Actually, keep a sensible default: if `--model` not provided, search `~/.cache/huggingface/hub/` for any `models--*Qwen*` directory. If exactly one found, use it. Otherwise, print error asking for `--model`.
|
| 440 |
+
|
| 441 |
+
- [ ] **Step 2.3: Remove the old #define block (lines 72-140)**
|
| 442 |
+
|
| 443 |
+
Delete the entire block of `#define`s from `HIDDEN_DIM` through `MODEL_PATH_DEFAULT`. Also remove:
|
| 444 |
+
- `#define NUM_FULL_ATTN_LAYERS 10` (line ~1011)
|
| 445 |
+
- `#define NUM_LINEAR_LAYERS 30` (line ~1033)
|
| 446 |
+
|
| 447 |
+
- [ ] **Step 2.4: Remove old static array declarations**
|
| 448 |
+
|
| 449 |
+
Remove these declarations (they're now in `alloc_tracking_arrays()`):
|
| 450 |
+
- `static int g_expert_freq[NUM_LAYERS][NUM_EXPERTS];` (line 206)
|
| 451 |
+
- `static uint8_t g_expert_seen[NUM_LAYERS][NUM_EXPERTS / 8];` (line 214)
|
| 452 |
+
- `static LZ4IndexEntry *g_lz4_index[NUM_LAYERS];` (line 198)
|
| 453 |
+
- `static uint8_t g_cache_seen[NUM_LAYERS][NUM_EXPERTS];` (line 254)
|
| 454 |
+
- `static uint64_t g_cache_last_touch_token[NUM_LAYERS][NUM_EXPERTS];` (line 255)
|
| 455 |
+
- `static uint64_t g_cache_last_evict_token[NUM_LAYERS][NUM_EXPERTS];` (line 256)
|
| 456 |
+
- `static int g_pred_experts[60][MAX_K];` (line 3298)
|
| 457 |
+
- `static int g_pred_count[60];` (line 3299)
|
| 458 |
+
|
| 459 |
+
- [ ] **Step 2.5: Update active_expert_size() to use cfg**
|
| 460 |
+
|
| 461 |
+
Change:
|
| 462 |
+
```c
|
| 463 |
+
static inline size_t active_expert_size(void) {
|
| 464 |
+
return g_use_2bit ? EXPERT_SIZE_2BIT : EXPERT_SIZE;
|
| 465 |
+
}
|
| 466 |
+
```
|
| 467 |
+
to:
|
| 468 |
+
```c
|
| 469 |
+
static inline size_t active_expert_size(void) {
|
| 470 |
+
return g_use_2bit ? cfg.expert_size_2bit : cfg.expert_size_4bit;
|
| 471 |
+
}
|
| 472 |
+
```
|
| 473 |
+
|
| 474 |
+
- [ ] **Step 2.6: Attempt build — expect ~960 errors from removed defines**
|
| 475 |
+
|
| 476 |
+
```bash
|
| 477 |
+
cd metal_infer && make 2>&1 | head -50
|
| 478 |
+
```
|
| 479 |
+
|
| 480 |
+
This confirms the scope of replacements needed. Do NOT try to fix yet — just verify the errors are all "use of undeclared identifier" for the removed defines.
|
| 481 |
+
|
| 482 |
+
- [ ] **Step 2.7: Commit (broken state, WIP)**
|
| 483 |
+
|
| 484 |
+
```bash
|
| 485 |
+
git add metal_infer/infer.m
|
| 486 |
+
git commit -m "wip: remove old defines, wire up config loader (broken — refs not yet updated)"
|
| 487 |
+
```
|
| 488 |
+
|
| 489 |
+
---
|
| 490 |
+
|
| 491 |
+
## Task 3: Bulk replace #define references with cfg.field
|
| 492 |
+
|
| 493 |
+
**Files:**
|
| 494 |
+
- Modify: `metal_infer/infer.m` — ~960 occurrences across the entire file
|
| 495 |
+
|
| 496 |
+
This is the largest task. Use find-and-replace for each define→cfg mapping. Order matters: replace longer names first to avoid partial matches (e.g., `LINEAR_TOTAL_KEY` before `LINEAR_KEY_DIM`).
|
| 497 |
+
|
| 498 |
+
- [ ] **Step 3.1: Replace core model dimension defines**
|
| 499 |
+
|
| 500 |
+
Apply these replacements throughout the file (use replace-all):
|
| 501 |
+
|
| 502 |
+
| Old | New |
|
| 503 |
+
|-----|-----|
|
| 504 |
+
| `HIDDEN_DIM` | `cfg.hidden_dim` |
|
| 505 |
+
| `NUM_LAYERS` | `cfg.num_layers` |
|
| 506 |
+
| `NUM_ATTN_HEADS` | `cfg.num_attn_heads` |
|
| 507 |
+
| `NUM_KV_HEADS` | `cfg.num_kv_heads` |
|
| 508 |
+
| `HEAD_DIM` | `cfg.head_dim` |
|
| 509 |
+
| `VOCAB_SIZE` | `cfg.vocab_size` |
|
| 510 |
+
| `RMS_NORM_EPS` | `cfg.rms_norm_eps` |
|
| 511 |
+
|
| 512 |
+
**CAUTION:** Be careful with `HEAD_DIM` — it must not match inside longer names. Verify no `*_HEAD_DIM` defines exist that would be corrupted.
|
| 513 |
+
|
| 514 |
+
- [ ] **Step 3.2: Replace MoE defines**
|
| 515 |
+
|
| 516 |
+
| Old | New |
|
| 517 |
+
|-----|-----|
|
| 518 |
+
| `NUM_EXPERTS_PER_TOK` | `cfg.num_experts_per_tok` |
|
| 519 |
+
| `NUM_EXPERTS` | `cfg.num_experts` |
|
| 520 |
+
| `MOE_INTERMEDIATE` | `cfg.moe_intermediate` |
|
| 521 |
+
| `SHARED_INTERMEDIATE` | `cfg.shared_intermediate` |
|
| 522 |
+
| `GROUP_SIZE` | `cfg.group_size` |
|
| 523 |
+
| `BITS` | `cfg.bits` |
|
| 524 |
+
|
| 525 |
+
**CAUTION:** `NUM_EXPERTS` must not match `NUM_EXPERTS_PER_TOK`. Replace `NUM_EXPERTS_PER_TOK` first, then `NUM_EXPERTS`. Similarly, `BITS` is short — verify it doesn't appear in other contexts.
|
| 526 |
+
|
| 527 |
+
- [ ] **Step 3.3: Replace linear attention defines**
|
| 528 |
+
|
| 529 |
+
Replace in this order (longest first):
|
| 530 |
+
|
| 531 |
+
| Old | New |
|
| 532 |
+
|-----|-----|
|
| 533 |
+
| `LINEAR_TOTAL_VALUE` | `cfg.linear_total_value` |
|
| 534 |
+
| `LINEAR_TOTAL_KEY` | `cfg.linear_total_key` |
|
| 535 |
+
| `LINEAR_CONV_DIM` | `cfg.linear_conv_dim` |
|
| 536 |
+
| `LINEAR_NUM_V_HEADS` | `cfg.linear_num_v_heads` |
|
| 537 |
+
| `LINEAR_NUM_K_HEADS` | `cfg.linear_num_k_heads` |
|
| 538 |
+
| `LINEAR_KEY_DIM` | `cfg.linear_key_dim` |
|
| 539 |
+
| `LINEAR_VALUE_DIM` | `cfg.linear_value_dim` |
|
| 540 |
+
| `CONV_KERNEL_SIZE` | `cfg.conv_kernel_size` |
|
| 541 |
+
|
| 542 |
+
- [ ] **Step 3.4: Replace RoPE and full attention defines**
|
| 543 |
+
|
| 544 |
+
| Old | New |
|
| 545 |
+
|-----|-----|
|
| 546 |
+
| `FULL_ATTN_INTERVAL` | — (see Step 3.8 for formula replacement) |
|
| 547 |
+
| `ROPE_THETA` | `cfg.rope_theta` |
|
| 548 |
+
| `PARTIAL_ROTARY` | `cfg.partial_rotary` |
|
| 549 |
+
| `ROTARY_DIM` | `cfg.rotary_dim` |
|
| 550 |
+
|
| 551 |
+
- [ ] **Step 3.5: Replace expert offset defines**
|
| 552 |
+
|
| 553 |
+
4-bit offsets:
|
| 554 |
+
|
| 555 |
+
| Old | New |
|
| 556 |
+
|-----|-----|
|
| 557 |
+
| `EXPERT_SIZE` (but NOT `EXPERT_SIZE_2BIT`) | `cfg.expert_size_4bit` |
|
| 558 |
+
| `GATE_W_OFF_4` | `cfg.gate_w_off_4` |
|
| 559 |
+
| `GATE_S_OFF_4` | `cfg.gate_s_off_4` |
|
| 560 |
+
| `GATE_B_OFF_4` | `cfg.gate_b_off_4` |
|
| 561 |
+
| `UP_W_OFF_4` | `cfg.up_w_off_4` |
|
| 562 |
+
| `UP_S_OFF_4` | `cfg.up_s_off_4` |
|
| 563 |
+
| `UP_B_OFF_4` | `cfg.up_b_off_4` |
|
| 564 |
+
| `DOWN_W_OFF_4` | `cfg.down_w_off_4` |
|
| 565 |
+
| `DOWN_S_OFF_4` | `cfg.down_s_off_4` |
|
| 566 |
+
| `DOWN_B_OFF_4` | `cfg.down_b_off_4` |
|
| 567 |
+
|
| 568 |
+
2-bit offsets:
|
| 569 |
+
|
| 570 |
+
| Old | New |
|
| 571 |
+
|-----|-----|
|
| 572 |
+
| `EXPERT_SIZE_2BIT` | `cfg.expert_size_2bit` |
|
| 573 |
+
| `GATE_W_OFF_2` | `cfg.gate_w_off_2` |
|
| 574 |
+
| ... (same pattern for all 2-bit offsets) |
|
| 575 |
+
|
| 576 |
+
- [ ] **Step 3.6: Replace special token defines**
|
| 577 |
+
|
| 578 |
+
| Old | New |
|
| 579 |
+
|-----|-----|
|
| 580 |
+
| `EOS_TOKEN_1` | `cfg.eos_token_ids[0]` |
|
| 581 |
+
| `EOS_TOKEN_2` | `cfg.eos_token_ids[1]` |
|
| 582 |
+
| `THINK_START_TOKEN` | `cfg.think_start_token` |
|
| 583 |
+
| `THINK_END_TOKEN` | `cfg.think_end_token` |
|
| 584 |
+
| `MAX_SEQ_LEN` | `cfg.max_seq_len` |
|
| 585 |
+
|
| 586 |
+
- [ ] **Step 3.7: Replace struct-local defines**
|
| 587 |
+
|
| 588 |
+
| Old | New |
|
| 589 |
+
|-----|-----|
|
| 590 |
+
| `NUM_FULL_ATTN_LAYERS` | `cfg.num_full_attn_layers` |
|
| 591 |
+
| `NUM_LINEAR_LAYERS` | `cfg.num_linear_layers` |
|
| 592 |
+
|
| 593 |
+
- [ ] **Step 3.8: Replace FULL_ATTN_INTERVAL formula patterns**
|
| 594 |
+
|
| 595 |
+
Find all occurrences of the formula pattern `(i + 1) % FULL_ATTN_INTERVAL == 0` (or variants with different variable names like `layer_idx`) and replace with `cfg.is_full_attn[i]`.
|
| 596 |
+
|
| 597 |
+
Also replace index computation formulas:
|
| 598 |
+
- `(layer_idx + 1) / FULL_ATTN_INTERVAL - 1` → `cfg.full_attn_index[layer_idx]`
|
| 599 |
+
- `layer_idx - (layer_idx + 1) / FULL_ATTN_INTERVAL` → `cfg.linear_index[layer_idx]`
|
| 600 |
+
|
| 601 |
+
- [ ] **Step 3.9: Build and fix any remaining errors**
|
| 602 |
+
|
| 603 |
+
```bash
|
| 604 |
+
cd metal_infer && make 2>&1 | head -100
|
| 605 |
+
```
|
| 606 |
+
|
| 607 |
+
Fix any remaining compilation errors from the replacements. Common issues:
|
| 608 |
+
- `sizeof()` on old arrays that are now pointers (need explicit size)
|
| 609 |
+
- `memset()` on old arrays (need explicit size calculation)
|
| 610 |
+
- Places where `EXPERT_SIZE` was used generically (should use `active_expert_size()`)
|
| 611 |
+
|
| 612 |
+
- [ ] **Step 3.10: Commit**
|
| 613 |
+
|
| 614 |
+
```bash
|
| 615 |
+
git add metal_infer/infer.m
|
| 616 |
+
git commit -m "feat: replace all model #defines with cfg.* struct fields (~960 occurrences)"
|
| 617 |
+
```
|
| 618 |
+
|
| 619 |
+
---
|
| 620 |
+
|
| 621 |
+
## Task 4: Convert MetalCtx fixed arrays to dynamic allocation
|
| 622 |
+
|
| 623 |
+
**Files:**
|
| 624 |
+
- Modify: `metal_infer/infer.m` — MetalCtx struct (~line 1011) and metal_setup() (~line 1049)
|
| 625 |
+
|
| 626 |
+
- [ ] **Step 4.1: Change MetalCtx arrays to pointers**
|
| 627 |
+
|
| 628 |
+
In the `MetalCtx` struct, change:
|
| 629 |
+
```c
|
| 630 |
+
id<MTLBuffer> buf_kv_k[NUM_FULL_ATTN_LAYERS];
|
| 631 |
+
id<MTLBuffer> buf_kv_v[NUM_FULL_ATTN_LAYERS];
|
| 632 |
+
```
|
| 633 |
+
to:
|
| 634 |
+
```c
|
| 635 |
+
id<MTLBuffer> *buf_kv_k;
|
| 636 |
+
id<MTLBuffer> *buf_kv_v;
|
| 637 |
+
```
|
| 638 |
+
|
| 639 |
+
And:
|
| 640 |
+
```c
|
| 641 |
+
id<MTLBuffer> buf_delta_state[NUM_LINEAR_LAYERS];
|
| 642 |
+
id<MTLBuffer> buf_conv_state[NUM_LINEAR_LAYERS];
|
| 643 |
+
```
|
| 644 |
+
to:
|
| 645 |
+
```c
|
| 646 |
+
id<MTLBuffer> *buf_delta_state;
|
| 647 |
+
id<MTLBuffer> *buf_conv_state;
|
| 648 |
+
```
|
| 649 |
+
|
| 650 |
+
- [ ] **Step 4.2: Allocate in metal_setup()**
|
| 651 |
+
|
| 652 |
+
In `metal_setup()`, after creating the Metal device but before buffer allocation, add:
|
| 653 |
+
|
| 654 |
+
```c
|
| 655 |
+
ctx->buf_kv_k = calloc(cfg.num_full_attn_layers, sizeof(id<MTLBuffer>));
|
| 656 |
+
ctx->buf_kv_v = calloc(cfg.num_full_attn_layers, sizeof(id<MTLBuffer>));
|
| 657 |
+
ctx->buf_delta_state = calloc(cfg.num_linear_layers, sizeof(id<MTLBuffer>));
|
| 658 |
+
ctx->buf_conv_state = calloc(cfg.num_linear_layers, sizeof(id<MTLBuffer>));
|
| 659 |
+
```
|
| 660 |
+
|
| 661 |
+
- [ ] **Step 4.3: Convert stack VLAs in serve loop to malloc**
|
| 662 |
+
|
| 663 |
+
In the serve function (~line 6160), change:
|
| 664 |
+
```c
|
| 665 |
+
void *gpu_delta_snapshots[NUM_LINEAR_LAYERS];
|
| 666 |
+
void *gpu_conv_snapshots[NUM_LINEAR_LAYERS];
|
| 667 |
+
```
|
| 668 |
+
to:
|
| 669 |
+
```c
|
| 670 |
+
void **gpu_delta_snapshots = calloc(cfg.num_linear_layers, sizeof(void *));
|
| 671 |
+
void **gpu_conv_snapshots = calloc(cfg.num_linear_layers, sizeof(void *));
|
| 672 |
+
```
|
| 673 |
+
|
| 674 |
+
Add `free(gpu_delta_snapshots); free(gpu_conv_snapshots);` at function exit/cleanup.
|
| 675 |
+
|
| 676 |
+
- [ ] **Step 4.4: Convert main() stack arrays to use cfg**
|
| 677 |
+
|
| 678 |
+
In main() (~line 6801), change:
|
| 679 |
+
```c
|
| 680 |
+
int layer_fds[NUM_LAYERS];
|
| 681 |
+
int layer_fds_cold[NUM_LAYERS];
|
| 682 |
+
void *layer_mmaps[NUM_LAYERS];
|
| 683 |
+
size_t layer_mmap_sizes[NUM_LAYERS];
|
| 684 |
+
```
|
| 685 |
+
to VLAs using `cfg.num_layers` (C99 VLAs are fine here since main() runs after config load):
|
| 686 |
+
```c
|
| 687 |
+
int layer_fds[cfg.num_layers];
|
| 688 |
+
int layer_fds_cold[cfg.num_layers];
|
| 689 |
+
void *layer_mmaps[cfg.num_layers];
|
| 690 |
+
size_t layer_mmap_sizes[cfg.num_layers];
|
| 691 |
+
```
|
| 692 |
+
|
| 693 |
+
Or use malloc if compiler doesn't support VLAs in ObjC.
|
| 694 |
+
|
| 695 |
+
- [ ] **Step 4.5: Update LayerWeightCache array**
|
| 696 |
+
|
| 697 |
+
Change:
|
| 698 |
+
```c
|
| 699 |
+
static LayerWeightCache layer_cache[NUM_LAYERS];
|
| 700 |
+
```
|
| 701 |
+
to:
|
| 702 |
+
```c
|
| 703 |
+
static LayerWeightCache *layer_cache = NULL;
|
| 704 |
+
```
|
| 705 |
+
|
| 706 |
+
And in `alloc_tracking_arrays()` add:
|
| 707 |
+
```c
|
| 708 |
+
layer_cache = calloc(cfg.num_layers, sizeof(LayerWeightCache));
|
| 709 |
+
```
|
| 710 |
+
|
| 711 |
+
- [ ] **Step 4.6: Fix memset/sizeof on converted arrays**
|
| 712 |
+
|
| 713 |
+
Find all `memset(g_expert_seen, 0, sizeof(g_expert_seen))` and similar calls that relied on compile-time sizeof. Replace with explicit size:
|
| 714 |
+
```c
|
| 715 |
+
memset(g_expert_seen, 0, cfg.num_layers * ((cfg.num_experts + 7) / 8));
|
| 716 |
+
```
|
| 717 |
+
|
| 718 |
+
Same for `g_cache_seen`, `g_cache_last_touch_token`, `g_cache_last_evict_token`.
|
| 719 |
+
|
| 720 |
+
- [ ] **Step 4.7: Update expert_is_seen / expert_mark_seen helpers**
|
| 721 |
+
|
| 722 |
+
Change from direct 2D array access to flattened access using the helper macros:
|
| 723 |
+
```c
|
| 724 |
+
static inline int expert_is_seen(int layer, int expert) {
|
| 725 |
+
return (EXPERT_SEEN_BYTE(layer, expert) >> (expert & 7)) & 1;
|
| 726 |
+
}
|
| 727 |
+
static inline void expert_mark_seen(int layer, int expert) {
|
| 728 |
+
EXPERT_SEEN_BYTE(layer, expert) |= (1 << (expert & 7));
|
| 729 |
+
}
|
| 730 |
+
```
|
| 731 |
+
|
| 732 |
+
- [ ] **Step 4.8: Update g_pred_experts / g_pred_count access**
|
| 733 |
+
|
| 734 |
+
Replace all `g_pred_experts[layer_idx][k]` with `PRED_EXPERT(layer_idx, k)` and `g_pred_count[layer_idx]` with `PRED_COUNT(layer_idx)`.
|
| 735 |
+
|
| 736 |
+
- [ ] **Step 4.9: Update all g_expert_freq, g_cache_seen access patterns**
|
| 737 |
+
|
| 738 |
+
Replace all `g_expert_freq[layer][expert]` with `FREQ(layer, expert)` etc. throughout the file.
|
| 739 |
+
|
| 740 |
+
- [ ] **Step 4.10: Build and verify**
|
| 741 |
+
|
| 742 |
+
```bash
|
| 743 |
+
cd metal_infer && make clean && make
|
| 744 |
+
```
|
| 745 |
+
|
| 746 |
+
Expected: compiles with 0 errors, 0 warnings.
|
| 747 |
+
|
| 748 |
+
- [ ] **Step 4.11: Commit**
|
| 749 |
+
|
| 750 |
+
```bash
|
| 751 |
+
git add metal_infer/infer.m
|
| 752 |
+
git commit -m "feat: convert fixed-size arrays to dynamic allocation"
|
| 753 |
+
```
|
| 754 |
+
|
| 755 |
+
---
|
| 756 |
+
|
| 757 |
+
## Task 5: Validate with current model
|
| 758 |
+
|
| 759 |
+
**Files:** None (testing only)
|
| 760 |
+
|
| 761 |
+
- [ ] **Step 5.1: Run inference with 35B model**
|
| 762 |
+
|
| 763 |
+
```bash
|
| 764 |
+
cd metal_infer && ./infer --model ~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit --prompt "What is 2+2?" --tokens 20
|
| 765 |
+
```
|
| 766 |
+
|
| 767 |
+
Verify:
|
| 768 |
+
1. Config summary prints correctly to stderr
|
| 769 |
+
2. No NaN values
|
| 770 |
+
3. Coherent output
|
| 771 |
+
4. Same token/s as before (~4.7 tok/s)
|
| 772 |
+
|
| 773 |
+
- [ ] **Step 5.2: Run server mode**
|
| 774 |
+
|
| 775 |
+
```bash
|
| 776 |
+
cd metal_infer && ./infer --model ~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit --serve 8000
|
| 777 |
+
```
|
| 778 |
+
|
| 779 |
+
In another terminal:
|
| 780 |
+
```bash
|
| 781 |
+
curl -X POST http://localhost:8000/v1/chat/completions \
|
| 782 |
+
-H "Content-Type: application/json" \
|
| 783 |
+
-d '{"messages":[{"role":"user","content":"Hello"}]}'
|
| 784 |
+
```
|
| 785 |
+
|
| 786 |
+
Verify: no bus error, coherent response.
|
| 787 |
+
|
| 788 |
+
- [ ] **Step 5.3: Run with --timing flag**
|
| 789 |
+
|
| 790 |
+
```bash
|
| 791 |
+
cd metal_infer && ./infer --model ~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit --prompt "Hello" --tokens 10 --timing
|
| 792 |
+
```
|
| 793 |
+
|
| 794 |
+
Verify: timing breakdown looks normal, no regressions.
|
| 795 |
+
|
| 796 |
+
- [ ] **Step 5.4: Commit validation**
|
| 797 |
+
|
| 798 |
+
```bash
|
| 799 |
+
git add metal_infer/infer.m
|
| 800 |
+
git commit -m "feat: runtime model config — validated with Qwen3.5-35B-A3B"
|
| 801 |
+
```
|
| 802 |
+
|
| 803 |
+
---
|
| 804 |
+
|
| 805 |
+
## Task 6: Clean up and update header comment
|
| 806 |
+
|
| 807 |
+
**Files:**
|
| 808 |
+
- Modify: `metal_infer/infer.m:1-43` (header comment)
|
| 809 |
+
|
| 810 |
+
- [ ] **Step 6.1: Update file header comment**
|
| 811 |
+
|
| 812 |
+
Replace the hardcoded model description in the header (lines 1-43) with a generic description that mentions runtime config loading:
|
| 813 |
+
|
| 814 |
+
```c
|
| 815 |
+
/*
|
| 816 |
+
* infer.m — Qwen3.5 MoE inference engine using Metal
|
| 817 |
+
*
|
| 818 |
+
* Full forward pass: embedding -> N transformer layers -> norm -> lm_head -> sample
|
| 819 |
+
* Model architecture loaded at runtime from HuggingFace config.json (--model flag).
|
| 820 |
+
* Non-expert weights loaded from model_weights.bin (mmap'd at startup)
|
| 821 |
+
* Expert weights loaded from packed_experts/ per layer per token (pread)
|
| 822 |
+
*
|
| 823 |
+
* Supported models: Qwen3.5-35B-A3B, Qwen3.5-397B-A17B, and compatible MoE variants
|
| 824 |
+
* ...
|
| 825 |
+
```
|
| 826 |
+
|
| 827 |
+
- [ ] **Step 6.2: Update the startup banner in main()**
|
| 828 |
+
|
| 829 |
+
Change `printf("=== Qwen3.5-35B-A3B Metal Inference Engine ===\n")` to dynamically show the model info:
|
| 830 |
+
|
| 831 |
+
```c
|
| 832 |
+
printf("=== Flash-MoE Metal Inference Engine ===\n");
|
| 833 |
+
printf("Config: %s/config.json\n", cfg.model_path);
|
| 834 |
+
```
|
| 835 |
+
|
| 836 |
+
- [ ] **Step 6.3: Final build and test**
|
| 837 |
+
|
| 838 |
+
```bash
|
| 839 |
+
cd metal_infer && make clean && make && ./infer --prompt "Hello" --tokens 5
|
| 840 |
+
```
|
| 841 |
+
|
| 842 |
+
- [ ] **Step 6.4: Commit**
|
| 843 |
+
|
| 844 |
+
```bash
|
| 845 |
+
git add metal_infer/infer.m
|
| 846 |
+
git commit -m "chore: update header and banner for runtime model config"
|
| 847 |
+
```
|