<!--
  MLX-RESEARCH.md
  Zyquo Local

  Author: Simon-Pierre Boucher
  Mail: contact@spboucher.ai
-->

# MLX Swift Stack Research — Phase 0.A

Researched 2026-07-30 against live official sources (raw file fetches at exact
tags, GitHub releases API) plus empirical build tests on this machine
(M5 Max, 48 GB, macOS 27.0 beta, Swift 6.4 CLT).

## 1. Packages and SPM coordinates

| Package | Repo | Version used | Products we link |
|---|---|---|---|
| mlx-swift | `https://github.com/ml-explore/mlx-swift` | **0.31.6** (2026-07-02) | (transitive) MLX, MLXNN, MLXFast… |
| mlx-swift-lm | `https://github.com/ml-explore/mlx-swift-lm` | **3.31.4** (2026-06-30) | `MLXLLM`, `MLXLMCommon`, `MLXHuggingFace` |
| swift-huggingface | `https://github.com/huggingface/swift-huggingface` | **0.9.0** | `HuggingFace` |
| swift-transformers | `https://github.com/huggingface/swift-transformers` | **1.3.x** (1.3.3) | `Tokenizers` |

Key facts:

- The LLM layer (MLXLLM/MLXLMCommon/MLXVLM) **moved out of mlx-swift-examples**
  into `mlx-swift-lm`. mlx-swift-examples is frozen for libraries (last 2.29.1).
- **mlx-swift-lm 3.x is a breaking major**: it no longer depends on
  swift-transformers/Hub. `Downloader`, `Tokenizer`, `TokenizerLoader` are now
  protocols in MLXLMCommon; concrete impls come from the `MLXHuggingFace`
  macro library (`#huggingFaceTokenizerLoader()`, `#hubDownloader()`) backed by
  swift-huggingface + swift-transformers, which the app must declare itself.
- mlx-swift 0.31.5+ tags require a Swift ≥6.3 toolchain
  (`swift-tools-version: 6.3;(experimentalCGen)`). We have Swift 6.4 → fine.
- Minimum platform: macOS 14. Our `LSMinimumSystemVersion` = 14.0, arm64 only.

## 2. Loading, generating, streaming — verified API (tag 3.31.4)

Load from a **local directory** (our only load path — downloads are ours):

```swift
import MLXLLM, MLXLMCommon, MLXHuggingFace, Tokenizers

let container: ModelContainer = try await loadModelContainer(
    from: modelDirectoryURL,
    using: #huggingFaceTokenizerLoader()
)
```

- `LLMModelFactory._load` reads `config.json`, dispatches on `model_type` via
  `LLMTypeRegistry`, honors `generation_config.json` EOS overrides, and
  `loadWeights` enumerates **every `*.safetensors` recursively** (the shard
  index json is not consulted).
- There is **no progress callback for local-directory loads** — model load
  progress in UI is indeterminate (spinner + elapsed), not a percent bar.

Multi-turn with KV-cache reuse — `ChatSession` (MLXLMCommon):

```swift
let session = ChatSession(container,
    instructions: systemPrompt,
    generateParameters: params,
    history: restoredMessages)                       // rehydration supported
for try await g in session.streamDetails(to: prompt) {
    switch g {
    case .chunk(let text): ...
    case .info(let info):  ... // stats
    case .toolCall:        break
    }
}
session.clear()   // resets history + KV cache, keeps instructions
```

KV cache is held inside the session across turns (no re-prefill). ChatSession
is NOT thread-safe; `ModelContainer` is Sendable. One session per conversation,
owned by the `InferenceEngine` actor. Changing system prompt/params → rebuild
the session from persisted history (KV cache is lost, acceptable).

Lower level (available if needed): `container.prepare(input:)` +
`container.generate(input:parameters:) -> AsyncStream<Generation>`,
`generateTask(...)` for deterministic early-stop cleanup, `TokenIterator`.

`Generation` events: `.chunk(String)`, `.info(GenerateCompletionInfo)`,
`.toolCall(ToolCall)`. Stop reasons: `.stop`, `.length`, `.cancelled`.

### GenerateParameters (exact fields, verified)

`maxTokens: Int?`, `maxKVSize: Int?`, `kvBits: Int?` (4/8), `kvGroupSize`,
`quantizedKVStart`, `kvScheme: String?`, `temperature: Float` (default 0.6,
0 = greedy), `topP: Float` (1.0), `topK: Int` (0), `minP: Float` (0),
`seed: UInt64?`, `repetitionPenalty: Float?`, `repetitionContextSize` (20),
presence/frequency penalties + context sizes, `prefillStepSize` (512).

### Stats — GenerateCompletionInfo (final `.info` stream element)

`promptTokenCount`, `generationTokenCount`, `promptTime`, `generateTime`,
`stopReason`, computed `promptTokensPerSecond`, `tokensPerSecond`.
TTFT = time from send to first `.chunk` (we measure it ourselves).

### Unload

No explicit API: release all `ChatSession`/`ModelContainer` references, then
`MLX.GPU.clearCache()`. Verified guidance from the repo's own reference docs.

## 3. Model directory format (required files)

| File | Required | Notes |
|---|---|---|
| `config.json` | ✅ | `model_type`, quantization block, EOS ids |
| `*.safetensors` | ✅ | all shards, MLX-convention (scales/biases for quant) |
| `tokenizer.json` + `tokenizer_config.json` | ✅ | tokenizer + chat template |
| `model.safetensors.index.json` | ⬜ not read by Swift stack; keep anyway |
| `chat_template.jinja` / `*.jinja` | ⬜ newer HF layout |
| `generation_config.json` | ⬜ EOS override |
| `special_tokens_map.json` | ⬜ |

Download filter (mirrors the package's own): `*.safetensors`, `*.json`,
`*.jinja`. `DownloadManager` downloads exactly this set.

## 4. Supported `model_type` architectures (LLMTypeRegistry @ 3.31.4)

`mistral, mixtral, llama, phi, phi3, phimoe, gemma, gemma2, gemma3,
gemma3_text, gemma3n, gemma4, gemma4_unified, gemma4_text, qwen2, qwen3,
qwen3_moe, qwen3_next, qwen3_5, qwen3_5_moe, qwen3_5_text, minicpm,
starcoder2, cohere, openelm, internlm2, deepseek_v3, granite,
granitemoehybrid, mimo, mimo_v2_flash, minimax, glm4, glm4_moe,
glm4_moe_lite, acereason, falcon_h1, bitnet, smollm3, ernie4_5, lfm2,
baichuan_m1, exaone4, gpt_oss, lille-130m, olmoe, olmo2, olmo3, bailing_moe,
lfm2_moe, nanochat, nemotron_h, afmoe, jamba, mamba2, mistral3, apertus,
nemotron_labs_diffusion`

This list ships in the app (`SupportedArchitectures`) to warn before
downloading an unsupported model (checked against `config.model_type` from the
Hub API with `config=true`).

## 5. The no-Xcode build question — RESOLVED EMPIRICALLY

Upstream docs claim "SwiftPM (command line) cannot build the Metal shaders…
use xcodebuild". **This is outdated for Swift 6.4**: the new swift-build
backend runs `CompileMetalFile` on mlx-swift's `.metal` resources. Verified on
this machine:

1. CLT-only `swift build` fails with `unable to spawn process 'metal'` —
   the compiler is looked up **via `PATH`** (proven with a logging stub).
2. With `~/Developer/Metal.xctoolchain/usr/bin` (Metal Toolchain component
   v17.5.188, copied from cluster node M4M36) prepended to `PATH`,
   `swift build` succeeds and the binary executes on GPU (`Device(gpu, 0)`).
3. Most kernels are JIT-compiled at runtime by the Metal framework (no
   toolchain needed at runtime); AOT kernels (gemv, SDPA, norms, rope…) are
   compiled at build time into the `mlx-swift_Cmlx` resource bundle. Phase 2's
   PoC (real attention workload) re-confirms the metallib is bundled/found.

**Recipe** (full details + SwiftUI SDK pin in `docs/BUILD.md`): plain
`swift build` + Metal toolchain on PATH + `SDKROOT=MacOSX26.5.sdk`.
No Xcode IDE, no `.xcodeproj`, `xcodebuild` not needed.

For the `.app` bundle, the `mlx-swift_Cmlx.bundle` resource bundle must be
copied into `Contents/Resources/` (Makefile handles it); for bare CLI runs
from `.build/`, the bundle sits next to the executable so it is found.

## 6. Hub downloads: decision

**Custom `URLSession`-based `DownloadManager`** (per-file pause/resume/cancel,
HTTP Range resume across restarts, app-managed model folder), NOT
swift-huggingface's snapshot API (no per-file pause handle, imposes HF cache
layout). swift-transformers is still used for **tokenization** (non-negotiable
— it is what the MLX stack expects, via `#huggingFaceTokenizerLoader()`).
Full HF HTTP API contract documented in `docs/MODELS.md` §1.

## 7. Memory model

APIs (verified @ mlx-swift 0.31.6): `MLX.GPU.activeMemory / cacheMemory /
peakMemory / snapshot() / resetPeakMemory() / set(cacheLimit:) /
set(memoryLimit:relaxed:) / clearCache() / deviceInfo()` (architecture,
memorySize, maxRecommendedWorkingSetSize).

Sizing rules used by `MemoryAdvisor`:

- Weights ≈ `params × bits/8 × 1.08` (scales/biases overhead).
- KV cache ≈ `2 × layers × kvHeads × headDim × contextTokens × 2 bytes` (fp16);
  `kvBits` 8/4 halves/quarters it; `maxKVSize` caps it (rotating cache).
- +10–20 % working set for activations; macOS GPU-wired ceiling ≈ 70–75 % of
  unified RAM → verdicts: **Fits** if weights+KV ≤ 60 % of physical RAM,
  **Tight** ≤ 75 %, **Too large** above.
- Physical RAM via `sysctl hw.memsize`; live footprint via `GPU.snapshot()`.

## 8. Open flags

- mlx-swift #430 (metallib via build plugin) and #416 (`setMetallibPath`) were
  open as of 2026-07-30 — not relied upon.
- API surface verified at tag 3.31.4; `main` has drifted (MLXFoundationModels,
  MLXGuidedGeneration) — re-verify before any version bump.