From 5a79495c8a89b33230a6a8822064502052b1bc33 Mon Sep 17 00:00:00 2001 From: Aegis AI Assistant Date: Tue, 14 Apr 2026 14:24:54 -0700 Subject: [PATCH 1/3] feat: add upstream sync workflow, downstream CI and fork protection tests --- .github/workflows/downstream_integration.yml | 53 ++++++++++++++++ .github/workflows/upstream_sync.yml | 67 ++++++++++++++++++++ Tests/MLXTests/ForkProtectionTests.swift | 33 ++++++++++ UPSTREAM_MERGE_PLAN.md | 35 ++++++++++ 4 files changed, 188 insertions(+) create mode 100644 .github/workflows/downstream_integration.yml create mode 100644 .github/workflows/upstream_sync.yml create mode 100644 Tests/MLXTests/ForkProtectionTests.swift create mode 100644 UPSTREAM_MERGE_PLAN.md diff --git a/.github/workflows/downstream_integration.yml b/.github/workflows/downstream_integration.yml new file mode 100644 index 00000000..fe0ea7e4 --- /dev/null +++ b/.github/workflows/downstream_integration.yml @@ -0,0 +1,53 @@ +name: Downstream Integration Test + +on: + pull_request: + branches: [ main ] + +jobs: + test-swiftlm-umbrella: + name: Test with SwiftLM (Umbrella Repo) + runs-on: macos-15 + timeout-minutes: 40 + steps: + - name: Checkout PR branch (mlx-swift) + uses: actions/checkout@v4 + with: + fetch-depth: 0 + path: mlx-swift-pr + + - name: Checkout SwiftLM (Parent Repo) + uses: actions/checkout@v4 + with: + repository: SharpAI/SwiftLM + path: SwiftLM + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Install Metal Toolchain + run: xcodebuild -downloadComponent MetalToolchain || true + + - name: Resolve Dependencies (to populate SPM cache) + run: swift package resolve + working-directory: SwiftLM + + - name: Override mlx-swift dependency with PR checkout + # SPM `package edit` puts a local editable copy of the package in + # .build/checkouts; we then replace its contents with our PR branch. + run: | + swift package edit mlx-swift --path $GITHUB_WORKSPACE/mlx-swift-pr + working-directory: SwiftLM + + - name: Build SwiftLM (Release) + run: swift build -c release --build-tests + working-directory: SwiftLM + + - name: Install MLX Metal library + run: | + python3 -m venv /tmp/mlx_venv + /tmp/mlx_venv/bin/pip install --quiet mlx + cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib SwiftLM/.build/release/ + find SwiftLM/.build -type d -name "MLXLMTests" -exec cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib {}/ \; + + - name: Run SwiftLM Integration Tests + run: swift test -c release --parallel --skip-build + working-directory: SwiftLM diff --git a/.github/workflows/upstream_sync.yml b/.github/workflows/upstream_sync.yml new file mode 100644 index 00000000..9aa9aa43 --- /dev/null +++ b/.github/workflows/upstream_sync.yml @@ -0,0 +1,67 @@ +name: Upstream Sync + +on: + schedule: + - cron: '0 */4 * * *' # Run every 4 hours + workflow_dispatch: # Allow manual triggering + +jobs: + check-and-create-pr: + name: Check Upstream and Create PR + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + fetch-depth: 0 # Required to fetch all history for commit comparison + + - name: Check for upstream changes and set up branch + id: check_upstream + run: | + # Add the upstream repository + git remote add upstream https://github.com/ml-explore/mlx-swift.git + git fetch upstream main + + # Check how many commits upstream is ahead of our local main + AHEAD=$(git rev-list --count HEAD..upstream/main) + echo "Upstream is $AHEAD commits ahead." + + if [ "$AHEAD" -gt 0 ]; then + echo "has_changes=true" >> $GITHUB_OUTPUT + BRANCH_NAME="sync/upstream-latest" + echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT + + # Checkout upstream main into the static sync branch (force reset if it exists) + git checkout -B $BRANCH_NAME upstream/main + git push --force origin $BRANCH_NAME + else + echo "has_changes=false" >> $GITHUB_OUTPUT + echo "No new commits to sync." + fi + + - name: Create or Update Pull Request + if: steps.check_upstream.outputs.has_changes == 'true' + run: | + # Check if a PR already exists for this branch + PR_EXISTS=$(gh pr list HEAD --head ${{ steps.check_upstream.outputs.branch_name }} --json number --jq 'length') + + if [ "$PR_EXISTS" -eq "0" ]; then + gh pr create \ + --title "🔄 Auto-Sync: Upstream changes from ml-explore/mlx-swift" \ + --body "Automated PR to sync the latest changes from upstream \`ml-explore/mlx-swift\`'s \`main\` branch. + + ### Merge Warnings + * **Do NOT blind merge!** This repo contains custom SharpAI out-of-core SSD features. + * Refer to the \`UPSTREAM_MERGE_PLAN.md\` file locally to check standard SharpAI metal optimization discrepancies. + * Do **not** overwrite the submodule branch pointers under \`Package.swift\` without consultation." \ + --base main \ + --head ${{ steps.check_upstream.outputs.branch_name }} + else + echo "PR already exists. Force pushing to the branch updated the PR automatically." + fi + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/Tests/MLXTests/ForkProtectionTests.swift b/Tests/MLXTests/ForkProtectionTests.swift new file mode 100644 index 00000000..dd72e8e1 --- /dev/null +++ b/Tests/MLXTests/ForkProtectionTests.swift @@ -0,0 +1,33 @@ +import XCTest +import Foundation + +final class ForkProtectionTests: XCTestCase { + + /// This test ensures that the SharpAI custom SSD streaming kernel + /// (`moe_stream.metal`) is natively preserved in the generated metal code and sources. + /// If an upstream merge resets the metal codegen generator, this will safely fail. + func testSSDStreamerMetalCodegenExists() throws { + // Look for the metal source in the parent package structure to verify the submodule holds our patches + let fileManager = FileManager.default + let packageRoot = URL(fileURLWithPath: #file) + .deletingLastPathComponent() + .deletingLastPathComponent() + .deletingLastPathComponent() + + let customKernelPath = packageRoot + .appendingPathComponent("Source/Cmlx/mlx/mlx/backend/metal/kernels/moe_stream.metal") + + XCTAssertTrue(fileManager.fileExists(atPath: customKernelPath.path), + "🚨 ALARM: The SharpAI custom ssd-stream metal kernel is missing! A blind upstream Apple merge erased moe_stream.metal! Do not merge this PR.") + } + + /// This test verifies that the custom ThreadPool modification we added to load.cpp is still tracked. + func testThreadPoolFeatureExists() throws { + let fileManager = FileManager.default + let packageRoot = URL(fileURLWithPath: #file).deletingLastPathComponent().deletingLastPathComponent().deletingLastPathComponent() + + let threadPoolPath = packageRoot.appendingPathComponent("Source/Cmlx/mlx/mlx/threadpool.h") + XCTAssertTrue(fileManager.fileExists(atPath: threadPoolPath.path), + "🚨 ALARM: The SharpAI custom ThreadPool is missing! Did the upstream ml-explore merge drop the background loader?") + } +} diff --git a/UPSTREAM_MERGE_PLAN.md b/UPSTREAM_MERGE_PLAN.md new file mode 100644 index 00000000..36c6e68c --- /dev/null +++ b/UPSTREAM_MERGE_PLAN.md @@ -0,0 +1,35 @@ +# SharpAI `mlx-swift` Custom Patches & Upstream Merge Guide + +This document captures the divergence between the `SharpAI/mlx-swift` custom fork and the official `ml-explore/mlx-swift` upstream repository. **Do not overwrite these commits or configurations** during an upstream merge unless Apple officially merges the underlying features. + +These patches enable out-of-core model inference, specifically for SSD Flash Streaming, by relying on custom bindings and our dedicated `SharpAI/mlx` and `SharpAI/mlx-c` forks. + +### 1. Speculative Decoding +_Note: Core speculative decoding logic (Drafting, KV Cache routing) primarily exists natively in `SwiftBuddy`/`mlx-swift-lm`. Modifications inside `mlx-swift` strictly pertain to the underlying memory/cache alignment requirements (like bounds limit resolutions) rather than the decoding algorithms themselves._ + +### 2. SSD Streaming & Out-of-core Patches +Our core optimizations extend the native Engine bindings to allow pulling model matrices from disk to GPU memory directly. +- `6d3a11f` / `9e10176`: **SSD thread-pooling** - Implements background thread-pools in the wrappers with dynamic API toggles to manage off-main-thread I/O operations. +- `14f267b`: **GLM-5.1 compatibility limit fix** - Removes obsolete buffer bounds checks in the `ssd_streamer` that originally crashed multi-expert GLM instances. +- `8540be8` / `877d992`: **C-API extensions** - Restores modifications to the Swift bridge and `fast.h` to explicitly surface C++ streaming kernels (`moe_stream_op`). +- `b45b31c`: **Gemma 4 Apple Silicon array mapping** - Customized Metal kernels optimized for Gemma 4 memory configurations. + +### 3. Metal Code Generation & JIT Extensions +Because Swift package compilation automatically generates C++ metal headers, the JIT builder was intercepted to include our custom kernel `.metal` files. +- `61590a5`: Adds the custom `moe_stream.metal` processing logic directly into the codegen workflow. +- `7bb740b`: Tracks the SharpAI custom metal kernels alongside Apple originals so the JIT builder doesn't ignore them. +- `cbbf55c`: Enforces that the `fast.h` method signatures exactly match the upstreamed `fast.cpp` limits after code generation. +- `441ef64`: Adds `steel_conv_3d` C++ string targeting internally for the `Cmlx` build step. + +### 4. Build Configuration & Submodule Pointers +The `Package.swift` and internal git components have been locked to point to our ecosystem. Overwriting these with Apple's pointers will break the compiler. +- `2f60e7b`: Re-points Xcode and SPM to `SharpAI/mlx` and `SharpAI/mlx-c` dedicated GitHub forks instead of `ml-explore`. +- `72daf6c` / `d0d852f`: Retains specific submodule inline configurations holding metadata/device streams specifically built to support the SharpAI metal optimizations. +- `6139f94` / `41c8c4f`: Explicit `mlx-c` pointer bumps required to maintain compilation signatures for `fftshift` and `Shape`. +- `a4b0d27`: Manual override bumping the `cxxLanguageStandard` parameter directly to `.gnucxx20` to guarantee C++ compatibility parity. + +--- +### When merging upstream... +1. Execute `git fetch upstream` (where upstream is `https://github.com/ml-explore/mlx-swift.git`). +2. Do **not** use `--strategy-option=theirs` blindly! +3. Any conflicts in `Package.swift` (submodule references), `Cmlx/`, or `Metal/` codegen generators MUST favor the **SharpAI local implementation (HEAD)** to prevent kernel erasure. From a95188acf762eadb92e152697975eefbda3bd481 Mon Sep 17 00:00:00 2001 From: Aegis AI Assistant Date: Tue, 14 Apr 2026 15:57:08 -0700 Subject: [PATCH 2/3] fix(ci): use debug build to prevent test framework stripping --- .github/workflows/downstream_integration.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/downstream_integration.yml b/.github/workflows/downstream_integration.yml index fe0ea7e4..471e37d6 100644 --- a/.github/workflows/downstream_integration.yml +++ b/.github/workflows/downstream_integration.yml @@ -37,17 +37,17 @@ jobs: swift package edit mlx-swift --path $GITHUB_WORKSPACE/mlx-swift-pr working-directory: SwiftLM - - name: Build SwiftLM (Release) - run: swift build -c release --build-tests + - name: Build SwiftLM + run: swift build --build-tests working-directory: SwiftLM - name: Install MLX Metal library run: | python3 -m venv /tmp/mlx_venv /tmp/mlx_venv/bin/pip install --quiet mlx - cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib SwiftLM/.build/release/ + cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib SwiftLM/.build/debug/ || echo "No debug metal lib" find SwiftLM/.build -type d -name "MLXLMTests" -exec cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib {}/ \; - name: Run SwiftLM Integration Tests - run: swift test -c release --parallel --skip-build + run: swift test --parallel --skip-build working-directory: SwiftLM From 69c05cc7e9eae461e7095b7a0c4d011f68bd17ca Mon Sep 17 00:00:00 2001 From: Aegis AI Assistant Date: Tue, 14 Apr 2026 16:22:38 -0700 Subject: [PATCH 3/3] fix(ci): scatter metallib to all test binary locations --- .github/workflows/downstream_integration.yml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/downstream_integration.yml b/.github/workflows/downstream_integration.yml index 471e37d6..ec878c8a 100644 --- a/.github/workflows/downstream_integration.yml +++ b/.github/workflows/downstream_integration.yml @@ -45,8 +45,13 @@ jobs: run: | python3 -m venv /tmp/mlx_venv /tmp/mlx_venv/bin/pip install --quiet mlx - cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib SwiftLM/.build/debug/ || echo "No debug metal lib" - find SwiftLM/.build -type d -name "MLXLMTests" -exec cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib {}/ \; + METALLIB=/tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib + # Copy to all locations the test runner and linked bundles may search + cp $METALLIB SwiftLM/.build/debug/ || true + find SwiftLM/.build -name "*.xctest" -exec cp $METALLIB {}/Contents/MacOS/ \; 2>/dev/null || true + find SwiftLM/.build -type d \( -name "debug" -o -name "release" \) -exec cp $METALLIB {}/ \; 2>/dev/null || true + # Also copy beside test binaries directly + find SwiftLM/.build -name "*Tests" -type f -exec sh -c 'cp '$METALLIB' "$(dirname {})"/' \; 2>/dev/null || true - name: Run SwiftLM Integration Tests run: swift test --parallel --skip-build