feat: Integrates Upstream Sync & Fork Protection

.github/workflows/downstream_integration.yml

@@ -0,0 +1,58 @@
+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
+        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
+          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
+        working-directory: SwiftLM

.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 }}

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?")
+    }
+}

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.