extro.rook.host
Fork for thermals request add-json-schema-dpeq
4.5 kB
112 lines
1# vmlx-swift — local development fork
2
3This is a local working copy of `osaurus-ai/mlx-swift` (branch
4`osaurus-0.31.3`) set up as the upstream for `vmlx-swift-lm` optimization
5work. **Not pushed to any remote.** Edits here flow directly into
6`../vmlx-swift-lm` on the next `swift build` because `vmlx-swift-lm`'s
7`Package.swift` uses a local path dependency.
8
9## Layout
10
11| Path | Purpose |
12|------|---------|
13| `~/vmlx-swift` | this fork |
14| `~/vmlx-swift-lm` | model library consuming it |
15
16## Branches
17
18| Branch | Purpose |
19|--------|---------|
20| `osaurus-0.31.3` | pristine `osaurus-ai/mlx-swift` tracking branch, don't edit |
21| `vmlx-0.31.3` (HEAD) | working branch for optimizations |
22
23Submodule `Source/Cmlx/mlx` also has a local branch:
24
25| Branch | Purpose |
26|--------|---------|
27| (detached at upstream) | default `ml-explore/mlx` state |
28| `vmlx-patches-0.31.3` (HEAD) | committed patches for `GatherQMM::output_shapes` + `CustomKernel::set_output_shapes` that `vmlx-swift-lm`'s compiled-kernel path depends on |
29
30## Remotes
31
32| Name | URL | Purpose |
33|------|-----|---------|
34| `osaurus-upstream` | `https://github.com/osaurus-ai/mlx-swift.git` | pull latest osaurus work into `osaurus-0.31.3` |
35
36**No `origin` remote.** A plain `git push` will fail. This is intentional —
37the fork is local-only until we decide to publish.
38
39## Iteration workflow
40
411. Edit a file under `Source/MLX/…` or `Source/Cmlx/mlx/mlx/…`.
422. In `../vmlx-swift-lm`, run `swift build -c release`. SPM picks up the
43 local path dep automatically and rebuilds the Cmlx C++ + MLX Swift
44 modules against the new source.
453. Run a bench (`swift run -c release CompileBench` for op-level,
46 `swift run -c release RunBench` for full Qwen 3.5-35B multi-turn).
474. Iterate.
48
49## Runtime quirk: `default.metallib`
50
51`swift build` from the terminal does NOT compile the `.metal` kernel
52files in `Source/Cmlx/mlx-generated/metal/` into a `default.metallib`
53(only Xcode's build phase does that). As a workaround:
54
55```sh
56cp ~/Library/Developer/Xcode/DerivedData/vmlx-swift-lm-*/Build/Products/Release/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib \
57 ~/vmlx-swift-lm/default.metallib
58```
59
60The C++ loader falls through four paths before finally trying
61`default_mtllib_path` (= literal `"default.metallib"` resolved against
62cwd). Running `swift run` from `vmlx-swift-lm/` with a metallib sitting
63in that directory makes the load succeed.
64
65**TODO:** add an SPM build plugin or prebuild command to `Package.swift`
66that compiles `mlx-generated/metal/*.metal` into
67`Sources/Cmlx/Resources/default.metallib` and declares it as a resource.
68Then the Xcode DerivedData workaround goes away. This is blocked on
69deciding whether to patch Package.swift directly or submit the change
70upstream to `osaurus-ai/mlx-swift`.
71
72## Switching back to URL dependency
73
74When this fork is not needed, edit `vmlx-swift-lm/Package.swift`:
75
76```swift
77// From:
78.package(name: "mlx-swift", path: "../vmlx-swift"),
79// Back to:
80.package(url: "https://github.com/osaurus-ai/mlx-swift", branch: "osaurus-0.31.3"),
81```
82
83Then `swift package clean && swift build -c release`.
84
85## Current optimization targets
86
87Per `vmlx-swift-lm/docs/research/2026-04-13-decode-speed-to-120.md`, the
88Swift-vs-Python gap on Qwen 3.5-35B-A3B-4bit is ~10-12 tok/s end-to-end,
89and microbench attribution puts it at:
90
911. **~5 tok/s** — Swift eager per-op overhead (~0.5 us/op extra vs Python)
922. **~3 tok/s** — compile-site coverage differences (Python compiles
93 more than Swift captures)
943. **~3 tok/s** — structural: ARC on `MLXArray` class, module-boundary
95 dispatch, heap alloc on every op return
96
97Concrete things to try in this fork:
98
99- [ ] `@inlinable` on hot-path `MLXArray+Ops.swift` methods (`exp`,
100 `add`, `multiply`, `matmul`, `reshape`) **plus** `MLXArray.init(_:)`
101 so cross-module inlining actually works end-to-end
102- [ ] `@inline(__always)` on `MLXArray.ctx` getter to avoid dispatch on
103 every bridge call
104- [ ] Pool `mlx_array_new()` allocations to avoid repeated C-side
105 alloc/free on hot paths (requires mlx-c change)
106- [ ] Investigate whether `MLXArray` can become a struct (with a
107 `ManagedBuffer` holding `mlx_array`) instead of a `final class`.
108 High impact but breaks every caller.
109- [ ] Build plugin for Metal kernel compilation — removes the Xcode
110 DerivedData metallib workaround and enables clean CI builds
111
112Measurements go in the parent repo's research doc.