Tomas Marques commited on
Commit
2b41258
·
0 Parent(s):

Latest changes (squashed)

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .claude/settings.json +17 -0
  2. .gitignore +37 -0
  3. CLAUDE.md +445 -0
  4. FlashMoE-iOS/.gitignore +6 -0
  5. FlashMoE-iOS/397B_ANALYSIS.md +156 -0
  6. FlashMoE-iOS/App/AppDelegate.swift +22 -0
  7. FlashMoE-iOS/App/FlashMoEApp.swift +27 -0
  8. FlashMoE-iOS/Bridge/FlashMoE-Bridging-Header.h +5 -0
  9. FlashMoE-iOS/Bridge/FlashMoEBridge.swift +419 -0
  10. FlashMoE-iOS/FlashMoE.xcodeproj/project.pbxproj +516 -0
  11. FlashMoE-iOS/FlashMoEEngine/FlashMoEEngine.h +157 -0
  12. FlashMoE-iOS/FlashMoEEngine/FlashMoEEngine.m +1372 -0
  13. FlashMoE-iOS/IOS_PORT.md +347 -0
  14. FlashMoE-iOS/Models/ModelCatalog.swift +141 -0
  15. FlashMoE-iOS/REVIEW.md +206 -0
  16. FlashMoE-iOS/Resources/Assets.xcassets/AppIcon.appiconset/Contents.json +13 -0
  17. FlashMoE-iOS/Resources/Assets.xcassets/Contents.json +6 -0
  18. FlashMoE-iOS/Resources/FlashMoE.entitlements +16 -0
  19. FlashMoE-iOS/Services/DownloadManager.swift +512 -0
  20. FlashMoE-iOS/Views/BenchmarkView.swift +547 -0
  21. FlashMoE-iOS/Views/ChatView.swift +538 -0
  22. FlashMoE-iOS/Views/ContentView.swift +22 -0
  23. FlashMoE-iOS/Views/ModelDownloadRow.swift +232 -0
  24. FlashMoE-iOS/Views/ModelListView.swift +1212 -0
  25. FlashMoE-iOS/Views/ProfilerView.swift +274 -0
  26. FlashMoE-iOS/project.yml +75 -0
  27. README.md +1 -0
  28. app.py +164 -0
  29. autoresearch/baseline.txt +1 -0
  30. autoresearch/benchmark.sh +121 -0
  31. autoresearch/experiments.tsv +11 -0
  32. autoresearch/findings.md +35 -0
  33. autoresearch/prepare.sh +103 -0
  34. autoresearch/program.md +162 -0
  35. autoresearch/program_v2.md +212 -0
  36. autoresearch/research_queue.md +34 -0
  37. autoresearch/seed_papers.md +79 -0
  38. build_expert_index.py +167 -0
  39. cloud_repack.py +1355 -0
  40. copy_model_to_iphone.sh +140 -0
  41. docs/context-optimization.md +137 -0
  42. docs/expert-settings-guide.md +125 -0
  43. docs/io-and-gpu-exploration.md +215 -0
  44. docs/ios-port.md +244 -0
  45. docs/oom-prevention.md +136 -0
  46. docs/optimization-experiments-q4.md +155 -0
  47. docs/plan-async-pread-pipeline.md +166 -0
  48. docs/plan-io-experiments.md +198 -0
  49. docs/quantization-guide.md +187 -0
  50. 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
+ ![Progress](progress.png)
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
+ ```