diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index ef772ee88..5810c9d97 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -49,18 +49,26 @@ jobs: cd provider-swift swift build -c debug - - name: Extract mlx.metallib + - name: Build and stage source-matched mlx.metallib run: | - python3 -m venv /tmp/mlxvenv - /tmp/mlxvenv/bin/pip install 'mlx==0.31.2' - cp /tmp/mlxvenv/lib/python*/site-packages/mlx/lib/mlx.metallib \ - provider-swift/.build/debug/mlx.metallib + set -euo pipefail + command -v cmake >/dev/null 2>&1 || brew install cmake + METALLIB_CACHE_DIR="$RUNNER_TEMP/metallib-cache" \ + ./scripts/fetch-metallib.sh "$RUNNER_TEMP" + metallib="$RUNNER_TEMP/mlx.metallib" + provider_bin=$(cd provider-swift && swift build -c debug --show-bin-path) + for destination in provider-swift/.build/debug "$provider_bin"; do + mkdir -p "$destination" + cp "$metallib" "$destination/mlx.metallib" + done - name: Install Hugging Face client # Benchmarks serve the CBv2 default (gpt-oss-20b, ~12 GB) plus the # secondary multi-model checkpoint (gemma-4-26B QAT, ~14.5 GB) for # TestBenchmark_MultiModelMultiProvider. Warm HF cache ⇒ no-op. - run: /tmp/mlxvenv/bin/pip install huggingface_hub + run: | + python3 -m venv /tmp/mlxvenv + /tmp/mlxvenv/bin/pip install huggingface_hub - name: Download public models (pull request) run: | diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index de10cf91d..d0c41dab0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -120,9 +120,14 @@ jobs: ~/Library/org.swift.swiftpm provider-swift/.build libs/mlx-swift-lm/.build - key: spm-v2-${{ runner.os }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} + key: spm-v3-${{ runner.os }}-${{ github.sha }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} restore-keys: | - spm-v2-${{ runner.os }}- + spm-v3-${{ runner.os }}- + - name: Discard metallibs restored by the generic SwiftPM cache + run: | + for root in provider-swift/.build libs/mlx-swift-lm/.build; do + [ ! -d "$root" ] || find "$root" -type f -name mlx.metallib -delete + done - name: Install pinned Rust toolchain for production prompt parity run: | rustup toolchain install 1.88.0 --profile minimal @@ -136,28 +141,29 @@ jobs: # The root helper is a separately packaged executable product and is # not guaranteed to be linked merely because the test bundle builds. swift build --product darkbloom-fan-helper - - name: Extract mlx.metallib - working-directory: provider-swift + - name: Build and stage source-matched mlx.metallib run: | - python3 -m venv /tmp/mlxvenv - /tmp/mlxvenv/bin/pip install 'mlx==0.31.2' - metallib="/tmp/mlxvenv/lib/python$(/tmp/mlxvenv/bin/python -c 'import sys;print(f"{sys.version_info.major}.{sys.version_info.minor}")')/site-packages/mlx/lib/mlx.metallib" - # MLX's C++ loader (device.cpp load_default_library) searches for the - # metallib COLOCATED WITH THE RUNNING BINARY. Under `swift test` the - # running binary is the xctest runner INSIDE the .xctest bundle, so a - # copy in .build/debug alone is not found — it must sit next to the - # runner. Place it in both locations to cover `swift run`/built-binary - # paths and the test runner. Fail loudly if the source isn't there. - test -f "$metallib" || { echo "::error::mlx.metallib not found at $metallib"; exit 1; } - cp "$metallib" "$RUNNER_TEMP/mlx.metallib" - cp "$metallib" .build/debug/mlx.metallib - for bundle in .build/debug/*PackageTests.xctest; do - macos="$bundle/Contents/MacOS" - if [ -d "$macos" ]; then - cp "$metallib" "$macos/mlx.metallib" - echo "placed metallib in $macos" - fi + set -euo pipefail + command -v cmake >/dev/null 2>&1 || brew install cmake + METALLIB_CACHE_DIR="$RUNNER_TEMP/metallib-cache" \ + ./scripts/fetch-metallib.sh "$RUNNER_TEMP" + metallib="$RUNNER_TEMP/mlx.metallib" + test -s "$metallib" || { echo "::error::source-matched mlx.metallib was not produced"; exit 1; } + + # MLX loads the metallib from beside the running executable. Stage the + # helper-authorized artifact into every provider binary and XCTest + # location that the preceding build created. + provider_bin=$(cd provider-swift && swift build -c debug --show-bin-path) + for destination in provider-swift/.build/debug "$provider_bin"; do + mkdir -p "$destination" + cp "$metallib" "$destination/mlx.metallib" done + while IFS= read -r macos; do + cp "$metallib" "$macos/mlx.metallib" + echo "placed source-matched metallib in $macos" + done < <( + find provider-swift/.build -type d -path '*.xctest/Contents/MacOS' -print + ) - name: Run Swift tests working-directory: provider-swift # NOTE: do NOT pass --build-tests-relinking flags; the bundle is already @@ -186,7 +192,12 @@ jobs: # 12-vcpu runner); warm runs restore libs/mlx-swift-lm/.build from # the SwiftPM cache and finish in ~1 min. timeout-minutes: 35 - run: swift build --build-tests + run: | + # This standalone package otherwise resolves its remote mlx-swift + # branch while loading the root helper's local-source metallib. + swift package unedit --force mlx-swift >/dev/null 2>&1 || true + swift package edit --path ../mlx-swift mlx-swift + swift build --build-tests - name: Place metallib for nested test bundles id: place-nested-metallib if: ${{ !cancelled() && steps.build-nested-tests.outcome == 'success' }} @@ -293,9 +304,14 @@ jobs: ~/Library/org.swift.swiftpm provider-swift/.build libs/mlx-swift-lm/.build - key: spm-v2-${{ runner.os }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} + key: spm-v3-${{ runner.os }}-${{ github.sha }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} restore-keys: | - spm-v2-${{ runner.os }}- + spm-v3-${{ runner.os }}- + - name: Discard metallibs restored by the generic SwiftPM cache + run: | + for root in provider-swift/.build libs/mlx-swift-lm/.build; do + [ ! -d "$root" ] || find "$root" -type f -name mlx.metallib -delete + done - name: Build provider-swift working-directory: provider-swift run: | @@ -307,7 +323,12 @@ jobs: # parity, backend, eligibility, KV-sharing parity) from this nested # .build; saving it here keeps that lane warm across PRs. timeout-minutes: 35 - run: swift build --build-tests + run: | + # Keep the warmed nested workspace on the same local MLX source used + # by provider builds and scripts/fetch-metallib.sh. + swift package unedit --force mlx-swift >/dev/null 2>&1 || true + swift package edit --path ../mlx-swift mlx-swift + swift build --build-tests - name: Save SwiftPM cache uses: actions/cache/save@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0 with: @@ -316,7 +337,7 @@ jobs: ~/Library/org.swift.swiftpm provider-swift/.build libs/mlx-swift-lm/.build - key: spm-v2-${{ runner.os }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} + key: spm-v3-${{ runner.os }}-${{ github.sha }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} lint-console: name: Console UI Lint & Build diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 6d11e36b5..5f63fe7e7 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -46,12 +46,24 @@ jobs: - name: Build prompt-contract sidecar run: cargo build --manifest-path coordinator/promptsidecar/Cargo.toml --release - - name: Extract mlx.metallib + - name: Build and stage source-matched mlx.metallib run: | - python3 -m venv /tmp/mlxvenv - /tmp/mlxvenv/bin/pip install 'mlx==0.31.2' - cp /tmp/mlxvenv/lib/python*/site-packages/mlx/lib/mlx.metallib \ - provider-swift/.build/debug/mlx.metallib + set -euo pipefail + command -v cmake >/dev/null 2>&1 || brew install cmake + METALLIB_CACHE_DIR="$RUNNER_TEMP/metallib-cache" \ + ./scripts/fetch-metallib.sh "$RUNNER_TEMP" + metallib="$RUNNER_TEMP/mlx.metallib" + provider_bin=$(cd provider-swift && swift build -c debug --show-bin-path) + for destination in provider-swift/.build/debug "$provider_bin"; do + mkdir -p "$destination" + cp "$metallib" "$destination/mlx.metallib" + done + while IFS= read -r macos; do + cp "$metallib" "$macos/mlx.metallib" + done < <( + find provider-swift/.build libs/mlx-swift-lm/.build \ + -type d -path '*.xctest/Contents/MacOS' -print 2>/dev/null + ) - name: Install Hugging Face client # v0.7.5 one-engine: the testbed default is the smallest @@ -60,7 +72,9 @@ jobs: # register). Exact-cache routing additionally uses the small Gemma 4 # checkpoint because GPT-OSS's dynamic-date template is cold-only. # A warm HF cache makes these downloads no-ops. - run: /tmp/mlxvenv/bin/pip install huggingface_hub + run: | + python3 -m venv /tmp/mlxvenv + /tmp/mlxvenv/bin/pip install huggingface_hub - name: Download model (protected read-only auth) if: github.event_name == 'push' diff --git a/.github/workflows/release-swift.yml b/.github/workflows/release-swift.yml index c37311058..7dbb85671 100644 --- a/.github/workflows/release-swift.yml +++ b/.github/workflows/release-swift.yml @@ -236,9 +236,14 @@ jobs: ~/Library/org.swift.swiftpm provider-swift/.build libs/mlx-swift-lm/.build - key: spm-v2-${{ runner.os }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} + key: spm-v3-${{ runner.os }}-${{ github.sha }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} restore-keys: | - spm-v2-${{ runner.os }}- + spm-v3-${{ runner.os }}- + - name: Discard metallibs restored by the generic SwiftPM cache + run: | + for root in provider-swift/.build libs/mlx-swift-lm/.build; do + [ ! -d "$root" ] || find "$root" -type f -name mlx.metallib -delete + done - name: Verify production prompt parity run: | @@ -246,90 +251,62 @@ jobs: rustup override set 1.88.0 ./scripts/verify-prompt-parity.sh - # ---------------------------------------------------------------------- - # Build mlx.metallib (compiled Metal GPU kernels) FROM SOURCE. - # - # mlx-swift's Cmlx target does not compile .metal kernels via SwiftPM, so - # we build them here with cmake from the exact fork commit at - # libs/mlx-swift/Source/Cmlx/mlx (the same source the host C++ links - # against). This guarantees the kernels match — including the _nax kernels - # — and needs no published mlx wheel. The result is cached on the mlx - # commit + toolchain, so it only rebuilds when the submodule or Xcode - # changes (~1 min on a cold build, instant on a cache hit). - # ---------------------------------------------------------------------- - - - name: Resolve metallib cache key + # The root helper is the sole metallib builder. The Actions cache stores + # its private cache directory; the helper still runs on every release and + # authorizes the exact source/toolchain/deployment/JIT contract itself. + - name: Resolve source-matched metallib cache namespace id: mlxkey run: | set -euo pipefail MLX_SRC="libs/mlx-swift/Source/Cmlx/mlx" test -f "$MLX_SRC/mlx/version.h" || { echo "::error::mlx submodule not checked out at $MLX_SRC"; exit 1; } + SOURCE_STATUS=$(git -C "$MLX_SRC" status --porcelain=v1 --untracked-files=all) + if [ -n "$SOURCE_STATUS" ]; then + echo "::error::release metallib source checkout is dirty" + printf '%s\n' "$SOURCE_STATUS" + exit 1 + fi MLX_SHA=$(git -C "$MLX_SRC" rev-parse HEAD) + HELPER_SHA=$(shasum -a 256 scripts/fetch-metallib.sh | cut -d' ' -f1) SDK=$(xcrun --sdk macosx --show-sdk-version) - XCODE=$(xcodebuild -version | head -1 | tr -d ' ') - KEY="metallib-v1-${MLX_SHA}-${XCODE}-sdk${SDK}-dt${MLX_METALLIB_DEPLOYMENT_TARGET}-jitoff" - echo "mlx_src=$MLX_SRC" >> "$GITHUB_OUTPUT" - echo "mlx_sha=$MLX_SHA" >> "$GITHUB_OUTPUT" - echo "mlx_sha_short=${MLX_SHA:0:12}" >> "$GITHUB_OUTPUT" + XCODE=$(xcodebuild -version | tr -cs '[:alnum:].' '-') + KEY="metallib-helper-v1-${MLX_SHA}-${HELPER_SHA}-${XCODE}-sdk${SDK}-dt${MLX_METALLIB_DEPLOYMENT_TARGET}-jitoff" echo "key=$KEY" >> "$GITHUB_OUTPUT" - echo "Metallib cache key: $KEY" + echo "mlx_sha_short=${MLX_SHA:0:12}" >> "$GITHUB_OUTPUT" + echo "Metallib cache namespace: $KEY" - - name: Restore metallib cache + - name: Restore source-matched metallib cache id: metallib-cache uses: actions/cache/restore@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0 with: - path: ${{ runner.temp }}/metallib/mlx.metallib + path: ${{ runner.temp }}/metallib-cache key: ${{ steps.mlxkey.outputs.key }} - - name: Build mlx.metallib from source - if: steps.metallib-cache.outputs.cache-hit != 'true' + - name: Build source-matched mlx.metallib through root helper + id: metallib + env: + METALLIB_CACHE_DIR: ${{ runner.temp }}/metallib-cache run: | set -euo pipefail command -v cmake >/dev/null 2>&1 || brew install cmake - SRC="${{ steps.mlxkey.outputs.mlx_src }}" - BUILD="$RUNNER_TEMP/metallib-build" - OUT="$RUNNER_TEMP/metallib" - rm -rf "$BUILD"; mkdir -p "$OUT" - # JIT=OFF compiles the full kernel set into the metallib. Deployment - # target >= 26.2 + SDK >= 26.2 + Metal >= 4.0 is what gates the _nax - # kernels (mlx/backend/metal/kernels/CMakeLists.txt). - cmake -S "$SRC" -B "$BUILD" \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_OSX_DEPLOYMENT_TARGET="$MLX_METALLIB_DEPLOYMENT_TARGET" \ - -DMLX_METAL_JIT=OFF \ - -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_EXAMPLES=OFF \ - -DMLX_BUILD_BENCHMARKS=OFF -DMLX_BUILD_PYTHON_BINDINGS=OFF - cmake --build "$BUILD" --target mlx-metallib -j"$(sysctl -n hw.ncpu)" - cp "$BUILD/mlx/backend/metal/kernels/mlx.metallib" "$OUT/mlx.metallib" - ls -lh "$OUT/mlx.metallib" - - - name: Save metallib cache + MLX_SRC="libs/mlx-swift/Source/Cmlx/mlx" + test -z "$(git -C "$MLX_SRC" status --porcelain=v1 --untracked-files=all)" + ./scripts/fetch-metallib.sh "$RUNNER_TEMP/metallib" + test -z "$(git -C "$MLX_SRC" status --porcelain=v1 --untracked-files=all)" || { + echo "::error::metallib helper modified the release source checkout" + exit 1 + } + MLIB="$RUNNER_TEMP/metallib/mlx.metallib" + test -s "$MLIB" || { echo "::error::metallib missing at $MLIB"; exit 1; } + echo "metallib=$MLIB" >> "$GITHUB_OUTPUT" + + - name: Save source-matched metallib cache if: steps.metallib-cache.outputs.cache-hit != 'true' uses: actions/cache/save@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0 with: - path: ${{ runner.temp }}/metallib/mlx.metallib + path: ${{ runner.temp }}/metallib-cache key: ${{ steps.mlxkey.outputs.key }} - - name: Verify metallib completeness (nax + gemv) - id: metallib - run: | - set -euo pipefail - MLIB="$RUNNER_TEMP/metallib/mlx.metallib" - test -s "$MLIB" || { echo "::error::metallib missing at $MLIB"; exit 1; } - # Count with `grep -c` (not `grep -q`): grep -q closes the pipe on the - # first match, which makes `strings` on this 150 MB+ file exit via - # SIGPIPE and trip `set -o pipefail` — a false negative. - NAX=$(strings "$MLIB" | grep -c "_nax" || true) - GEMV=$(strings "$MLIB" | grep -c "gemv" || true) - # _nax kernels are required for the M5 Neural Accelerator path. Their - # absence means the build used the wrong SDK/deployment target (see the - # gate in kernels/CMakeLists.txt). gemv guards against an incomplete - # (JIT) build. Fail loudly rather than ship a broken bundle. - [ "$NAX" -gt 0 ] || { echo "::error::metallib has no _nax kernels (SDK/deployment target < 26.2?)"; exit 1; } - [ "$GEMV" -gt 0 ] || { echo "::error::metallib missing gemv kernels (incomplete build?)"; exit 1; } - echo "metallib=$MLIB" >> "$GITHUB_OUTPUT" - echo "Metallib OK: size=$(ls -lh "$MLIB" | awk '{print $5}') nax_refs=$NAX gemv_refs=$GEMV" - # ---------------------------------------------------------------------- # Build (tests run by CI, not the release workflow). # ---------------------------------------------------------------------- @@ -359,7 +336,7 @@ jobs: ~/Library/org.swift.swiftpm provider-swift/.build libs/mlx-swift-lm/.build - key: spm-v2-${{ runner.os }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} + key: spm-v3-${{ runner.os }}-${{ github.sha }}-${{ hashFiles('provider-swift/Package.resolved', 'libs/mlx-swift/Package.swift', 'libs/mlx-swift-lm/Package.swift') }} # Unit tests are run by CI (ci.yml + integration.yml) on push to # master. Skipped here to avoid flaky live MLX inference tests @@ -515,13 +492,17 @@ jobs: scripts/stage-swiftpm-resource-bundles.sh \ "$BIN_DIR" "$APP" "$RESOURCE_MANIFEST" - # Package-real runtime gate. This invokes the staged executable, - # resolves the resource from the installed app layout, then - # JIT-compiles and executes GPT-OSS bulk-write, decode-part, and - # decode-merge variants. The v0.7.6 artifact fails here because - # its bundle is absent. - DARKBLOOM_NO_UPDATE_CHECK=1 \ - "$APP/Contents/MacOS/$CLI_NAME" runtime-smoke + # Package-real runtime gate. Before its first GPU access the staged + # child decodes the retained Gemma config, proves authoritative + # overwrite of all three low-level controls, snapshots the early R1 + # latch/AOT capability without arming counters, then runs the + # existing GPT-OSS paged-kernel shapes. The v0.7.6 artifact still + # fails here because its resource bundle is absent. + PRE_SIGN_SMOKE_OUTPUT=$(DARKBLOOM_NO_UPDATE_CHECK=1 \ + "$APP/Contents/MacOS/$CLI_NAME" runtime-smoke) + printf '%s\n' "$PRE_SIGN_SMOKE_OUTPUT" + printf '%s\n' "$PRE_SIGN_SMOKE_OUTPUT" \ + | grep -Fqx 'gemma-optimizations-runtime-smoke: ok' # Hardened-runtime sign each binary and resource, then the bundle. # mlx.metallib must be signed before the bundle or codesign rejects @@ -686,6 +667,7 @@ jobs: # with the coordinator and verified by provider self-update. cmp "$STAGE/bin/$CLI_NAME" "$APP/Contents/MacOS/$CLI_NAME" cmp "$STAGE/bin/$ENCLAVE_NAME" "$APP/Contents/MacOS/$ENCLAVE_NAME" + cmp "$STAGE/bin/mlx.metallib" "$APP/Contents/MacOS/mlx.metallib" tar czf /tmp/darkbloom-bundle-macos-arm64.tar.gz -C "${{ steps.bundle.outputs.stage }}" . tar tzf /tmp/darkbloom-bundle-macos-arm64.tar.gz | sort | tee /tmp/darkbloom-bundle-files.txt @@ -721,6 +703,29 @@ jobs: rm -rf "$SMOKE_ROOT" mkdir -p "$SMOKE_ROOT" tar xzf /tmp/darkbloom-bundle-macos-arm64.tar.gz -C "$SMOKE_ROOT" + FINAL_APP_METALLIB="$SMOKE_ROOT/Darkbloom.app/Contents/MacOS/mlx.metallib" + FINAL_FLAT_METALLIB="$SMOKE_ROOT/bin/mlx.metallib" + test -s "$FINAL_APP_METALLIB" + test -s "$FINAL_FLAT_METALLIB" + cmp "$FINAL_FLAT_METALLIB" "$FINAL_APP_METALLIB" + # Repeat the helper's exact four-part completeness contract against + # the final extracted, signed bytes rather than the pre-signing input. + NAX_SYMBOL="_nax" + GEMV_SYMBOL="gemv" + R1_BUILDER_SYMBOL="build_gemma4_sorted_expert_tiles_bm32" + R1_KERNEL_SYMBOL="affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_alN_true_bm_32_bn_32_bk_32" + for symbol in \ + "$NAX_SYMBOL" \ + "$GEMV_SYMBOL" \ + "$R1_BUILDER_SYMBOL" \ + "$R1_KERNEL_SYMBOL" + do + MATCHES=$(strings "$FINAL_FLAT_METALLIB" | grep -F -c "$symbol" || true) + if [ "$MATCHES" -eq 0 ]; then + echo "::error::final signed metallib is missing required symbol: $symbol" + exit 1 + fi + done FINAL_FAN_HELPER="$SMOKE_ROOT/Darkbloom.app/Contents/Helpers/$FAN_HELPER_NAME" FINAL_FAN_MARKER="$SMOKE_ROOT/Darkbloom.app/Contents/Resources/darkbloom-runtime-capabilities/fan-helper-v1" test -f "$FINAL_FAN_HELPER" @@ -737,8 +742,11 @@ jobs: "-R=$FAN_HELPER_REQUIREMENT" "$FINAL_FAN_HELPER" codesign --verify --deep --strict --verbose=2 \ "$SMOKE_ROOT/Darkbloom.app" - DARKBLOOM_NO_UPDATE_CHECK=1 \ - "$SMOKE_ROOT/Darkbloom.app/Contents/MacOS/$CLI_NAME" runtime-smoke + FINAL_SMOKE_OUTPUT=$(DARKBLOOM_NO_UPDATE_CHECK=1 \ + "$SMOKE_ROOT/Darkbloom.app/Contents/MacOS/$CLI_NAME" runtime-smoke) + printf '%s\n' "$FINAL_SMOKE_OUTPUT" + printf '%s\n' "$FINAL_SMOKE_OUTPUT" \ + | grep -Fqx 'gemma-optimizations-runtime-smoke: ok' REPORTED=$("$SMOKE_ROOT/Darkbloom.app/Contents/MacOS/$CLI_NAME" --version) ./scripts/check-release-version.sh "$VERSION" "$REPORTED" test "$(/usr/libexec/PlistBuddy -c 'Print :CFBundleVersion' \ @@ -746,9 +754,9 @@ jobs: test "$(/usr/libexec/PlistBuddy -c 'Print :CFBundleShortVersionString' \ "$SMOKE_ROOT/Darkbloom.app/Contents/Info.plist")" = "$VERSION" - BINARY_HASH=$(shasum -a 256 "${{ steps.bundle.outputs.stage }}/bin/$CLI_NAME" | cut -d' ' -f1) + BINARY_HASH=$(shasum -a 256 "$SMOKE_ROOT/bin/$CLI_NAME" | cut -d' ' -f1) BUNDLE_HASH=$(shasum -a 256 /tmp/darkbloom-bundle-macos-arm64.tar.gz | cut -d' ' -f1) - METALLIB_HASH=$(shasum -a 256 "${{ steps.bundle.outputs.stage }}/bin/mlx.metallib" | cut -d' ' -f1) + METALLIB_HASH=$(shasum -a 256 "$FINAL_FLAT_METALLIB" | cut -d' ' -f1) echo "BINARY_HASH=$BINARY_HASH" >> "$GITHUB_ENV" echo "BUNDLE_HASH=$BUNDLE_HASH" >> "$GITHUB_ENV" echo "METALLIB_HASH=$METALLIB_HASH" >> "$GITHUB_ENV" diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ef9355f0..40990ba59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## Unreleased (v0.8.2 candidate — provider) + +### Provider (Swift) + +#### Performance + +- **Gemma 4 26B-A4B v0.8.2 optimization stack** — Layer-18 lazy prefill submission; coupled weighted-expert-unsort + safe-R1 expert-QMM gate (both default-on via `[gemma_optimizations]`); the VLM wrapper's directly shared text tower; packed multimodal prefill inside q=128 query blocks; source-matched metallib enforced across CI/release/packaged smoke. Final performance and retention deltas are pending a same-tree A/B measurement on the reviewed release tree. Earlier gitignored measurements predated the final kernel edits and are not release evidence. Dropped before the final cut: expert gate/up packing, dense gate/up packing, standalone weighted-unsort, standalone R1. `0cc5fc9c9` + +--- + ## Unreleased (Apr 26 - May 25, 2026) 26 commits since `aa74499`. diff --git a/Makefile b/Makefile index 846870b06..8240abfb2 100644 --- a/Makefile +++ b/Makefile @@ -44,11 +44,34 @@ prompt-sidecar: prompt-sidecar-format prompt-sidecar-check prompt-sidecar-test p # ---- Provider (Swift, Apple Silicon) -------------------------------------- -provider-build: ## swift build for the Swift provider CLI +provider-build: ## Build the Swift provider CLI with its source-matched metallib cd provider-swift && swift build - -provider-test: ## swift test for the Swift provider CLI - cd provider-swift && swift test + @set -eu; \ + bin_path="$$(cd provider-swift && swift build --show-bin-path)"; \ + ./scripts/fetch-metallib.sh "$$bin_path" + +provider-test: ## Build and run Swift provider tests with source-matched metallibs + cd provider-swift && swift build --build-tests + @set -eu; \ + bin_path="$$(cd provider-swift && swift build --show-bin-path)"; \ + ./scripts/fetch-metallib.sh "$$bin_path"; \ + runner_tmp=""; \ + trap 'test -z "$$runner_tmp" || rm -f "$$runner_tmp"' EXIT; \ + trap 'exit 143' HUP INT TERM; \ + found=0; \ + for bundle in "$$bin_path"/*PackageTests.xctest; do \ + [ -d "$$bundle" ] || continue; \ + runner_dir="$$bundle/Contents/MacOS"; \ + mkdir -p "$$runner_dir"; \ + runner_tmp="$$runner_dir/.mlx.metallib.$$$$"; \ + cp "$$bin_path/mlx.metallib" "$$runner_tmp"; \ + mv -f "$$runner_tmp" "$$runner_dir/mlx.metallib"; \ + runner_tmp=""; \ + found=1; \ + done; \ + [ "$$found" -eq 1 ] || { echo "provider test runner bundle not found in $$bin_path" >&2; exit 1; }; \ + trap - EXIT HUP INT TERM + cd provider-swift && swift test --skip-build provider: provider-build provider-test ## Build + test provider diff --git a/README.md b/README.md index 382436b1b..1453ee635 100644 --- a/README.md +++ b/README.md @@ -271,6 +271,10 @@ idle_timeout_mins = 60 # unload an idle model after N minutes (0 = never) max_model_slots = 3 # max models resident at once continuous_batching = true +[gemma_optimizations] +prefill_layer18 = true # default ON; 18-layer prefill submissions +weighted_r1 = true # default ON; coupled weighted-unsort + safe R1 + [coordinator] url = "wss://api.darkbloom.dev/ws/provider" heartbeat_interval_secs = 5 @@ -286,6 +290,11 @@ start = "22:00" end = "08:00" ``` +`[gemma_optimizations]` is authoritative for production provider processes. +Both selected controls default to `true` even when an older config omits the +section. To roll back either optimization, set its key to `false` and run +`darkbloom restart`; restart is the activation boundary. + ## Self-route & direct mode Running a node also makes your **own** inference free. diff --git a/coordinator/api/server.go b/coordinator/api/server.go index f6fb1b466..a6cf06797 100644 --- a/coordinator/api/server.go +++ b/coordinator/api/server.go @@ -151,7 +151,7 @@ func keyLimitResetFromContext(ctx context.Context) string { // the provider's EngineV2Factory.prepareProductionBackend for the argument). // Keep this fallback in sync with ProviderCore.version so dev/in-memory // coordinators advertise the same floor as the Swift binary they expect. -var LatestProviderVersion = "0.8.1" +var LatestProviderVersion = "0.8.2" // minProviderVersionForDesiredModels is the first provider version whose Swift // runtime understands the desired_models message. The coordinator must NOT send diff --git a/coordinator/api/toolschema.go b/coordinator/api/toolschema.go index 07aeff3a8..8dc8441c5 100644 --- a/coordinator/api/toolschema.go +++ b/coordinator/api/toolschema.go @@ -711,36 +711,6 @@ func assertionFamilyTypes(dict map[string]any) map[string]struct{} { return families } -// typelessAssertionFamiliesAmbiguous reports whether a typeless node's -// assertions cannot determine a single renderable type: either its assertion -// keywords span more than one type family (e.g. `{"minimum":5,"minLength":2}`) -// or its only assertion is `not` (e.g. `{"not":{"type":"string"}}`, which -// accepts every non-string — no injected type can preserve that, and the -// string default would make the schema unsatisfiable). Callers reject these -// before normalization; nodes with structural, union, or finite-value -// evidence are typed by those higher-priority rules instead. -func typelessAssertionFamiliesAmbiguous(dict map[string]any) bool { - for _, key := range []string{ - "properties", "patternProperties", "additionalProperties", - "items", "prefixItems", - } { - if _, ok := dict[key]; ok { - return false - } - } - if _, ok := unionMemberType(dict); ok { - return false - } - families := assertionFamilyTypes(dict) - if len(families) > 1 { - return true - } - if _, negated := dict["not"]; negated && len(families) == 0 { - return true - } - return false -} - // finiteValueTypes reports the JSON type names of a node's const/enum values: // the set of concrete (non-null) member types plus whether null appears. // ok is false when the node carries no const and no non-empty enum array. diff --git a/docs/provider/beta-features.md b/docs/provider/beta-features.md index 6ce588c47..a3ef58446 100644 --- a/docs/provider/beta-features.md +++ b/docs/provider/beta-features.md @@ -1,22 +1,48 @@ # Beta Features -Beta features are experimental provider capabilities that are **off by default**. -They are validated enough to try in production but may change, carry caveats, or -only apply to specific model families. Enable them per provider when you want to -opt in. +Beta features are experimental provider capabilities with feature-specific +defaults. The benchmark-selected Gemma stack is **on by default**, including for +existing configs that omit its section; reserved and opt-in features remain off +until enabled. Every feature can be changed per provider. ## How beta features are toggled -Beta features are **config-backed**, not environment-variable backed. Each one is -a field in your provider TOML config (`~/.config/darkbloom/provider.toml`). This -matters: the launchd daemon started by `darkbloom start` only inherits a tiny -allowlist of `DARKBLOOM_*` environment variables -(`provider-swift/Sources/ProviderCore/Daemon/LaunchAgent.swift`), so an env-var -toggle would silently have no effect on the running daemon. A TOML field is always -read by every serve path (daemon, `--foreground`, and `--local`). +Beta features are **config-backed**. Each one is a field in your provider TOML +config (`~/.config/darkbloom/provider.toml`), and that config is authoritative +for daemon, `--foreground`, and `--local` processes. Low-level environment +variables are implementation details, not an independent production control +surface. Process-wide optimization state is resolved at startup, so restart +after changing a restart-required feature. The registry of available features lives in -`provider-swift/Sources/ProviderCore/Config/BetaFeatures.swift`. +`BetaFeatures.all` in +`provider-swift/Sources/ProviderCore/Config/BetaFeatures.swift:73-133`. + +### Canonical implementation references + +- Default-on decoding for omitted keys is implemented by + `GemmaOptimizationSettings.init(from:)` in + `provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift:29-35`. + The missing-section fallback is `ProviderConfig.init(from:)` in + `provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift:391-400`. +- Startup projection before the first MLX device access (daemon, foreground, + `--local`, and `benchmark` processes): the shared ordering seam is + `provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift:24-35` + (`ServeRuntimePreparer.prepareRuntime`); the serve path calls it through + `Start.prepareServeRuntime` at + `provider-swift/Sources/darkbloom/StartCommand.swift:84-91,128-147`. + `Benchmark.run` rejects a conflicting shell override at + `provider-swift/Sources/darkbloom/BenchmarkCommand.swift:147-159`, prepares + the runtime at `:160-165`, and reports the effective controls at `:166-170`. + The guard helper is `ServeRuntimePreparer.conflictingEnvironmentOverride` + in `provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift:48-79`. +- The config→environment projection and its overwrite authority: + `GemmaOptimizationEnvironment.projection(for:)` in + `provider-swift/Sources/ProviderCore/Config/GemmaOptimizationEnvironment.swift:10-23` + and `GemmaOptimizationEnvironment.apply(_:)` at `:52-64`. +- The `benchmark`-selected coupling of weighted unsort with safe R1 as one + control: `GemmaOptimizationSettings.weightedR1` in + `provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift:10-14`. ## The `darkbloom beta` command @@ -28,19 +54,64 @@ darkbloom beta disable # turn a feature off ``` `enable`/`disable` perform a read-modify-write of the TOML config and print -whether a restart is required. Most beta features take effect on the next backend -start: +whether a restart is required. For example, to roll back the default-on coupled +Gemma expert optimization: ```bash -darkbloom beta enable mtp +darkbloom beta disable gemma-weighted-r1 darkbloom restart ``` -You can also see which beta features are active in `darkbloom status` (the -`Beta features:` line), or edit the TOML directly. +The restart is the activation boundary. Re-enable and restart to restore the +selected default. You can also see which beta features are active in +`darkbloom status` (the `Beta features:` line), or edit the TOML directly. ## Available features +### `gemma-prefill-layer18` — layer-18 prefill submission + +This optimization submits queued Gemma prefill work every 18 transformer +layers instead of waiting for one final submission. It defaults ON for both new +and pre-existing provider configs. + +```toml +[gemma_optimizations] +prefill_layer18 = true +``` + +Rollback is config-backed and restart-required: + +```bash +darkbloom beta disable gemma-prefill-layer18 +darkbloom restart +``` + +Setting the key to `false` (or using the command above) restores the legacy +one-final-submission behavior after restart. + +### `gemma-weighted-r1` — coupled weighted unsort + safe R1 + +This optimization defaults ON and is deliberately one atomic production +control. It enables both the direct weighted expert reduction and the safe +exact-shape R1 QMM path. There is no supported config or beta combination that +enables one without the other. + +```toml +[gemma_optimizations] +weighted_r1 = true +``` + +To roll back both paths together: + +```bash +darkbloom beta disable gemma-weighted-r1 +darkbloom restart +``` + +Missing `[gemma_optimizations]` sections and missing keys decode as `true`, so +old configs receive the selected v0.8.2 stack. An explicit `false` plus restart +is the durable rollback. + ### `mtp` — Gemma 4 multi-token prediction code path Providers v0.7.12 and later contain the default-off MTP implementation, but @@ -51,10 +122,14 @@ darkbloom beta enable mtp darkbloom restart ``` -Activation additionally requires a verified `spec_dec` assistant artifact in -the model catalog. Production publishes that artifact for -`gemma-4-26b-qat-4bit`; other models and providers without the explicit beta -setting continue with target-only decoding. +Without a valid local `[backend] mtp_drafter_path` override, activation also +requires a verified `spec_dec` assistant artifact in the model catalog. The +current public production catalog publishes `metadata.spec_dec` for +`gemma-4-26b-qat-4bit` ([live catalog](https://api.darkbloom.dev/v1/models/catalog?type=text)); +other catalog models and providers without the explicit beta setting continue +with target-only decoding. Artifact resolution and target-only fallback are in +`ProviderLoop.specDecPreparation` at +`provider-swift/Sources/ProviderCore/ProviderLoop+MTP.swift:35-54`. The implementation has target-authoritative verification and focused parity coverage. That is not a universal certification of token-identical behavior on diff --git a/docs/provider/cli-reference.md b/docs/provider/cli-reference.md index fc629c896..cfc2cf342 100644 --- a/docs/provider/cli-reference.md +++ b/docs/provider/cli-reference.md @@ -37,8 +37,10 @@ darkbloom start [flags] | `--bind ` | Bind address for local modes (default 127.0.0.1) | | `--no-auth` | Disable local API-key auth (trusted/airgapped only) | -Preflight checks (SIP, debugger, GPU, memory) run before the model picker -(`provider-swift/Sources/darkbloom/StartCommand.swift:429-448`). +Preflight checks for boot security, debugger attachment, and memory run before +the model picker (`provider-swift/Sources/darkbloom/StartCommand+Preflight.swift:9-27`); +the Metal requirement is enforced by `Start.prepareServeRuntime` +(`provider-swift/Sources/darkbloom/StartCommand.swift:128-147`). Examples: @@ -154,9 +156,10 @@ Two of the detailed checks cover the KV-backend rollout: | `kv backend posture` | An EXPLICIT `paged` or `contiguous` request was not honoured: refused (no engine built, the box serves nothing for that model) or silently degraded to another backend. | `auto` never fails this check — it promises nothing, so whichever backend it -lands on is honoured by definition. It resolves paged as of v0.8.0 and -degrades to contiguous on a box that cannot serve paged, so an `auto` slot -reporting contiguous is expected output, not a finding. When +lands on is honoured by definition. It resolves contiguous as of v0.8.1, so an +`auto` slot reporting contiguous is expected output, not a finding. Explicit +`paged` remains available and refuses a load it cannot serve instead of +silently changing backends. When the state file is past the wedge bar the backend verdict is WITHHELD rather than asserted from a snapshot that may predate a reload. @@ -258,9 +261,15 @@ This toggles `provider.auto_update` in `provider.toml`. ## `darkbloom beta` -Manage opt-in beta features. Beta features are off by default and config-backed -(a TOML field), so they apply to the launchd daemon too — unlike environment -variables, which the daemon does not inherit. +Manage configurable beta features. Defaults are feature-specific: the selected +Gemma optimizations default on, while reserved/opt-in features default off. +Provider TOML is authoritative for every serve mode. The Gemma defaults and +missing-key decode are defined in +`provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift:16-34`, +with the missing-section fallback in +`provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift:397-400`. +The shared pre-Metal projection is +`provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift:24-35`. ```bash darkbloom beta list # all features + on/off (default subcommand) @@ -271,12 +280,20 @@ darkbloom beta disable # turn off | Feature | Effect | |---------|--------| -| `mtp` | Default-off Gemma 4 MTP code path; requires a separately published and verified `spec_dec` artifact, which production does not currently have | +| `gemma-prefill-layer18` | Default-on layer-18 prefill submission; disable and restart for legacy submission behavior | +| `gemma-weighted-r1` | Default-on atomic weighted-unsort + safe-R1 pair; disable and restart to roll back both | +| `mtp` | Default-off Gemma 4 MTP code path; uses a valid local `mtp_drafter_path` or a verified catalog `spec_dec` artifact. The current production catalog publishes one for `gemma-4-26b-qat-4bit` | `enable`/`disable` read-modify-write the TOML config and report whether a restart -is required. See [Beta Features](beta-features.md) for the full guide. `darkbloom -beta list` also accepts `--json`. Installing a provider release does not enable -MTP, and local parity results are not a blanket M1–M3/unknown-chip certification. +is required. Restart is the activation boundary for process-wide optimization +state. The durable locked write and restart instruction are implemented in +`provider-swift/Sources/darkbloom/BetaCommand.swift:201-235`. See +[Beta Features](beta-features.md) for the full guide. `darkbloom beta list` also +accepts `--json`. Installing a provider release does not enable MTP, and local +parity results are not a blanket M1-M3/unknown-chip certification. +The published assistant metadata is visible in the +[public production catalog](https://api.darkbloom.dev/v1/models/catalog?type=text) +under `gemma-4-26b-qat-4bit.metadata.spec_dec`. `kv-quant` was removed in v0.8.0 and is no longer a valid feature id. ## `darkbloom fan` (experimental) diff --git a/docs/provider/hardware-requirements.md b/docs/provider/hardware-requirements.md index 8df3afe33..f17a2b45a 100644 --- a/docs/provider/hardware-requirements.md +++ b/docs/provider/hardware-requirements.md @@ -9,8 +9,8 @@ the models you want to serve. | Component | Minimum | Notes | |-----------|---------|-------| | **CPU** | Apple M1 (or later) | Apple Silicon required; Intel Macs are not supported | -| **RAM** | 8 GB | Start path rejects `< 8 GB` (`provider-swift/Sources/darkbloom/StartCommand.swift:444-447`) | -| **GPU** | Apple Silicon integrated GPU | CPU-only execution is rejected (`provider-swift/Sources/darkbloom/StartCommand.swift:80-85`) | +| **RAM** | 8 GB | Start path rejects `< 8 GB` (`provider-swift/Sources/darkbloom/StartCommand+Preflight.swift:23-27`) | +| **GPU** | Apple Silicon integrated GPU | CPU-only execution is rejected (`provider-swift/Sources/darkbloom/StartCommand.swift:128-147`) | | **Storage** | 50 GB free | SSD required; model weights are large | | **macOS** | 14 (Sonoma) | Newer is better; install script enforces Darwin + arm64 | | **Network** | Outbound HTTPS to coordinator | No inbound port is required | diff --git a/docs/provider/installation.md b/docs/provider/installation.md index 3ba9993d8..c9e6ec708 100644 --- a/docs/provider/installation.md +++ b/docs/provider/installation.md @@ -76,7 +76,7 @@ also treat warnings as failures. The canonical path is `~/.config/darkbloom/provider.toml`. The loader also reads legacy paths for backward compatibility; see -`provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift:214-252`. +`provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift:485-513`. A minimal example: @@ -92,11 +92,28 @@ enabled_models = [] idle_timeout_mins = 60 max_model_slots = 3 +[gemma_optimizations] +prefill_layer18 = true +weighted_r1 = true + [coordinator] url = "wss://api.darkbloom.dev/ws/provider" private_only = false ``` +Both Gemma controls default ON when the section or either key is absent, so +older configs receive the selected stack. Provider TOML is authoritative; set +a key to `false` and run `darkbloom restart` for a durable rollback. The defaults +and missing-key decode are canonical in +`provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift:16-34`, +with the missing-section fallback in +`provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift:397-400`. +Startup applies that config before Metal initialization +(`provider-swift/Sources/darkbloom/StartCommand.swift:84-91` and +`provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift:24-35`), and the +beta command's locked read-modify-write plus restart instruction is implemented +at `provider-swift/Sources/darkbloom/BetaCommand.swift:201-235`. + ## Updating the provider ```bash diff --git a/docs/provider/quickstart.md b/docs/provider/quickstart.md index 40d592b6a..15cf40ef0 100644 --- a/docs/provider/quickstart.md +++ b/docs/provider/quickstart.md @@ -15,10 +15,10 @@ serving the public fleet, or use the same node for your own free inference via | **Network** | Outbound HTTPS (port 443) | Low-latency path to `api.darkbloom.dev` | The installer enforces macOS + Apple Silicon up front -(`scripts/install.sh:41-48`). The start path rejects CPU-only execution via -`GPUEnforcement.requireMetal()` (`provider-swift/Sources/darkbloom/StartCommand.swift:80-85`) -and rejects machines with less than 8 GB RAM -(`provider-swift/Sources/darkbloom/StartCommand.swift:444-447`). +(`scripts/install.sh:41-48`). The start path rejects CPU-only execution through +`Start.prepareServeRuntime` (`provider-swift/Sources/darkbloom/StartCommand.swift:128-147`) +and rejects machines with less than 8 GB RAM in `Start.runPreflightChecks` +(`provider-swift/Sources/darkbloom/StartCommand+Preflight.swift:14-27`). ## Install @@ -109,7 +109,7 @@ It is created automatically on first start. The TOML schema is defined in `provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift`. ```toml -config_version = 1 +config_version = 2 [provider] name = "darkbloom-mac16-1" @@ -123,6 +123,10 @@ enabled_models = [] idle_timeout_mins = 60 max_model_slots = 3 +[gemma_optimizations] +prefill_layer18 = true +weighted_r1 = true + [coordinator] url = "wss://api.darkbloom.dev/ws/provider" heartbeat_interval_secs = 5 @@ -134,6 +138,26 @@ start = "22:00" end = "08:00" ``` +- `gemma_optimizations.prefill_layer18` — default ON, including when an older + config omits the section or key. Set to `false` and restart to restore legacy + one-final-submission Gemma prefill. The default and missing-key decode are in + `provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift:16-34`; + the missing-section fallback is in + `provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift:397-400`. +- `gemma_optimizations.weighted_r1` — default ON, including when omitted. This + is one atomic production control for weighted unsort and safe R1; the two + paths cannot be configured independently + (`provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift:10-18`, + coupled projection at + `provider-swift/Sources/ProviderCore/Config/GemmaOptimizationEnvironment.swift:14-22`). +- Provider TOML is authoritative for both controls. Changes take effect at + process restart; after setting either key to `false`, run `darkbloom restart` + to activate the rollback. The start path projects config before Metal access + (`provider-swift/Sources/darkbloom/StartCommand.swift:84-91` and + `provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift:24-35`), while + `darkbloom beta` durably locks, reloads, and saves the selected value before + printing the restart boundary + (`provider-swift/Sources/darkbloom/BetaCommand.swift:201-235`). - `backend.enabled_models` — if non-empty, only these models are advertised. - `backend.idle_timeout_mins` — minutes of inactivity before an idle model is unloaded (default 60; 0 disables eviction). @@ -175,14 +199,15 @@ end = "08:00" cache is not constructed there. Vision (VLM) models are NOT forced to contiguous. The VLM veto in `EngineV2KVBackendPolicy.applySlotVetoes` - (`guard isVLM, !pagedHonorsSpanMasks`, `provider-swift/Sources/ProviderCore/Inference/EngineV2KVBackendPolicy.swift:162`) + (`guard isVLM, !pagedHonorsSpanMasks`, `provider-swift/Sources/ProviderCore/Inference/EngineV2KVBackendPolicy.swift:202-210`) fires only when the paged cache does not affirm multimodal span masks, and `PagedLayerCache.honorsSpanMaskContextsByConstruction` is `true` - (`libs/mlx-swift-lm/Libraries/MLXLMCommon/ContinuousBatchingV2/Paged/PagedLayerCache.swift:982`), + (`libs/mlx-swift-lm/Libraries/MLXLMCommon/ContinuousBatchingV2/Paged/PagedLayerCache.swift:994`), which is what the slot factory passes - (`provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory.swift:190`), - so the veto is inert: a VLM slot gets paged under `"auto"` like any - other model. + (`provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory.swift:301-304`), + so the veto is inert: an explicitly paged VLM slot can use paged like any + other model. Under `"auto"`, every slot resolves contiguous as described + above. The concurrency cap above matters: paged only overtakes contiguous above ~5 concurrent rows, so pairing `engine_v2_kv_backend = "paged"` with a low `engine_v2_max_concurrent` @@ -192,14 +217,14 @@ end = "08:00" under an explicit `"paged"` the model REFUSES to load instead, with the underlying reason attached: `EngineV2KVBackendPolicy.degradesPagedFailure` - (`selection != .paged`, `provider-swift/Sources/ProviderCore/Inference/EngineV2KVBackendPolicy.swift:183`) + (`selection != .paged`, `provider-swift/Sources/ProviderCore/Inference/EngineV2KVBackendPolicy.swift:229-233`) returns `false` for — and only for — an explicit `.paged` selection, so a paged fleet can never silently serve contiguous. That refusal surfaces as a 503 and the coordinator reroutes: the engine-construction catch wraps it as `InferenceError.modelLoadFailed` (`provider-swift/Sources/ProviderCore/ProviderLoop+ModelLoading.swift:543-549`), `loadErrorStatusCode` maps that case to 503 - (`same file:975-983`), and the coordinator counts a 503 as + (`same file:979-1007`), and the coordinator counts a 503 as `capacityRejection` — no reputation strike — then cools the load-rejecting pair so retries skip it (`coordinator/api/provider.go:2332`, cool-down at `:2343-2351`). diff --git a/docs/provider/troubleshooting.md b/docs/provider/troubleshooting.md index 047325f30..a5b714830 100644 --- a/docs/provider/troubleshooting.md +++ b/docs/provider/troubleshooting.md @@ -34,8 +34,10 @@ is stale. Re-download from the coordinator and verify again. ### `darkbloom start` fails immediately -Preflight checks are in -`provider-swift/Sources/darkbloom/StartCommand.swift:429-448`. +Boot-security, debugger, and memory preflight checks are in +`provider-swift/Sources/darkbloom/StartCommand+Preflight.swift:9-27`; Metal +enforcement is in `Start.prepareServeRuntime` +(`provider-swift/Sources/darkbloom/StartCommand.swift:128-147`). | Error | Fix | |-------|-----| diff --git a/docs/releases/v0.8.0-notes.md b/docs/releases/v0.8.0-notes.md index 60c9d189e..d917096a5 100644 --- a/docs/releases/v0.8.0-notes.md +++ b/docs/releases/v0.8.0-notes.md @@ -561,4 +561,3 @@ production, investigation-only otherwise. no canary fleet, so the soak *is* the canary. Every measurable gate (G0a, G0b, G1, G2, G5) passes on real weights. Watch the first fleet day with `DARKBLOOM_CBV2_PAGED_KV=0` ready. - diff --git a/e2e/testbed/provider.go b/e2e/testbed/provider.go index 5e5de9026..bf3e2aef5 100644 --- a/e2e/testbed/provider.go +++ b/e2e/testbed/provider.go @@ -1,6 +1,7 @@ package testbed import ( + "bytes" "context" "fmt" "log/slog" @@ -8,6 +9,8 @@ import ( "os/exec" "path/filepath" "strings" + "syscall" + "time" ) func providerBuildConfig() string { @@ -44,34 +47,37 @@ func BuildProvider(ctx context.Context, logger *slog.Logger) (string, error) { providerDir := filepath.Join(repoRoot, "provider-swift") cfg := providerBuildConfig() - binaryPath := providerDir + "/.build/" + cfg + "/darkbloom" - if _, err := os.Stat(binaryPath); err == nil { - metallibPath := providerDir + "/.build/" + cfg + "/mlx.metallib" - if _, err2 := os.Stat(metallibPath); err2 == nil { - logger.Info("using cached provider binary", "path", binaryPath) - return binaryPath, nil - } + showBinPath := exec.CommandContext(ctx, "swift", "build", "-c", cfg, "--show-bin-path") + showBinPath.Dir = providerDir + binPathOutput, err := showBinPath.Output() + if err != nil { + return "", fmt.Errorf("resolve provider build path: %w", err) + } + binPath := strings.TrimSpace(string(binPathOutput)) + if binPath == "" { + return "", fmt.Errorf("resolve provider build path: swift returned an empty path") } + binaryPath := filepath.Join(binPath, "darkbloom") logger.Info("building provider binary", "dir", providerDir, "config", cfg) - cmd := exec.CommandContext(ctx, "swift", "build", "-c", cfg) cmd.Dir = providerDir - - out, err := cmd.CombinedOutput() - if err != nil { - return "", fmt.Errorf("swift build provider: %w: %s", err, string(out)) + out, buildErr := cmd.CombinedOutput() + if buildErr != nil { + return "", fmt.Errorf("swift build provider: %w: %s", buildErr, string(out)) } - - if _, err := os.Stat(binaryPath); err != nil { + if info, statErr := os.Stat(binaryPath); statErr != nil || info.IsDir() || info.Mode()&0o111 == 0 { return "", fmt.Errorf("provider binary not found after build: %s", binaryPath) } - if err := ensureMetallib(providerDir, logger); err != nil { + // Candidate binaries always receive a freshly staged metallib from the + // exact nested MLX source. An existing colocated file is not evidence that + // it matches the host code. + if err := ensureMetallib(ctx, repoRoot, binPath, logger); err != nil { return "", fmt.Errorf("metallib setup: %w", err) } - logger.Info("provider binary built", "path", binaryPath) + logger.Info("provider binary ready", "path", binaryPath) return binaryPath, nil } @@ -94,37 +100,69 @@ func findRepositoryRoot(start string) (string, error) { } } -func ensureMetallib(providerDir string, logger *slog.Logger) error { - cfg := providerBuildConfig() - metallibPath := providerDir + "/.build/" + cfg + "/mlx.metallib" - if _, err := os.Stat(metallibPath); err == nil { - return nil +func ensureMetallib( + ctx context.Context, + repoRoot string, + binPath string, + logger *slog.Logger, +) error { + helper := filepath.Join(repoRoot, "scripts", "fetch-metallib.sh") + info, err := os.Stat(helper) + if err != nil || info.IsDir() || info.Mode()&0o111 == 0 { + return fmt.Errorf("source metallib helper is not executable: %s", helper) } - if envPath := os.Getenv("MLX_METALLIB_PATH"); envPath != "" { - if _, err := os.Stat(envPath); err == nil { - return copyFile(envPath, metallibPath) - } + cmd := exec.Command(helper, binPath) + cmd.Dir = repoRoot + out, err := runProcessGroup(ctx, cmd) + if len(out) != 0 { + (&logWriter{logger: logger, prefix: "metallib helper"}).Write(out) } - - siteDirs, _ := filepath.Glob("/tmp/mlxvenv/lib/python*/site-packages/mlx/lib") - for _, dir := range siteDirs { - src := filepath.Join(dir, "mlx.metallib") - if _, err := os.Stat(src); err == nil { - logger.Info("copying mlx.metallib from Python wheel", "src", src) - return copyFile(src, metallibPath) - } + if err != nil { + return fmt.Errorf("build source-matched metallib: %w", err) } - return fmt.Errorf("mlx.metallib not found; install mlx==0.31.2 Python wheel and copy to %s or set MLX_METALLIB_PATH", metallibPath) + metallibPath := filepath.Join(binPath, "mlx.metallib") + metallib, err := os.Stat(metallibPath) + if err != nil || metallib.IsDir() || metallib.Size() == 0 { + return fmt.Errorf("source metallib helper did not stage %s", metallibPath) + } + return nil } -func copyFile(src, dst string) error { - data, err := os.ReadFile(src) - if err != nil { - return err +func runProcessGroup(ctx context.Context, cmd *exec.Cmd) ([]byte, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + var output bytes.Buffer + cmd.Stdout = &output + cmd.Stderr = &output + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + if err := cmd.Start(); err != nil { + return output.Bytes(), err + } + + done := make(chan error, 1) + go func() { + done <- cmd.Wait() + }() + + select { + case err := <-done: + return output.Bytes(), err + case <-ctx.Done(): + // Let the helper shell handle TERM and run its EXIT cleanup. If CMake + // or a compiler child does not exit promptly, kill the entire group. + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) + select { + case <-done: + case <-time.After(2 * time.Second): + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + <-done + } + return output.Bytes(), ctx.Err() } - return os.WriteFile(dst, data, 0644) } func findProviderBinary() string { diff --git a/libs/mlx-swift b/libs/mlx-swift index df1fdc5f7..97b3c9216 160000 --- a/libs/mlx-swift +++ b/libs/mlx-swift @@ -1 +1 @@ -Subproject commit df1fdc5f7821a1fabe921fdefbc42ac74dcfb6bc +Subproject commit 97b3c92162afab0ad7868ce364a40a2382c219a5 diff --git a/libs/mlx-swift-lm b/libs/mlx-swift-lm index 802398dc7..ed55bee83 160000 --- a/libs/mlx-swift-lm +++ b/libs/mlx-swift-lm @@ -1 +1 @@ -Subproject commit 802398dc7cedb607ad807f23cf37027e8fb49a2b +Subproject commit ed55bee83beb0623152f4c2e70f0cf99ad379e35 diff --git a/provider-swift/README.md b/provider-swift/README.md index acb6ce67e..d7ee50f92 100644 --- a/provider-swift/README.md +++ b/provider-swift/README.md @@ -12,11 +12,15 @@ This package is **CLI-only**: no SwiftUI app, no `.app` bundle, no DMG. ## Build & test ```bash -swift test -swift build -c release -# Outputs: -# .build/release/darkbloom -# .build/release/darkbloom-enclave +# From the repository root. These targets build first, stage the matching +# source-built metallib at every runtime path, and then run tests skip-build. +make provider-test +make provider-build + +# Optimized local binary: +cd provider-swift && swift build -c release && cd .. +./scripts/fetch-metallib.sh release +# Outputs include provider-swift/.build/release/{darkbloom,mlx.metallib}. ``` The package depends on local submodules at `../libs/mlx-swift` and `../libs/mlx-swift-lm`. Make sure they are checked out: @@ -30,19 +34,23 @@ git submodule update --init --recursive `mlx-swift`'s `Cmlx` target does **not** auto-compile its Metal kernels through SwiftPM. The runtime needs an `mlx.metallib` file colocated with the binary (or inside the binary's resource bundle), or it crashes on the first MLX call with `Failed to load the default metallib`. -Until we land a SwiftPM build-tool plugin for this, the workaround is to ship the matching `mlx.metallib` from the MLX Python wheel that pins the same C++ ABI as the fork (currently `mlx==0.31.2`). For local dev, the simplest setup is: +The canonical helper builds the metallib from the exact nested MLX source used +by `Cmlx` (`libs/mlx-swift/Source/Cmlx/mlx`), with Metal JIT disabled and the +required deployment target and kernel completeness checks. It never extracts a +Python wheel: ```bash -python3 -m venv /tmp/mlxvenv -/tmp/mlxvenv/bin/pip install 'mlx==0.31.2' -cp /tmp/mlxvenv/lib/python*/site-packages/mlx/lib/mlx.metallib \ - .build/release/mlx.metallib +# From the repository root, after the corresponding Swift build: +./scripts/fetch-metallib.sh debug +./scripts/fetch-metallib.sh release -# Then: -.build/release/darkbloom serve --foreground +provider-swift/.build/release/darkbloom serve --foreground ``` -`release-swift.yml` in CI does the same thing automatically and bakes the metallib into the released bundle next to `darkbloom`. +The helper also accepts an absolute destination directory and +`METALLIB_CACHE_DIR`; local, integration, CI, and release paths use this same +source-matched builder. Release packaging colocates the resulting metallib with +`darkbloom`. ## Layout @@ -110,6 +118,6 @@ Done in v0.5.0: Still pending: - [ ] Phase 4b: true continuous batching (deferrable past cutover; today's `BatchScheduler` does prefill-serial + decode-concurrent on a single ModelContainer). -- [ ] Phase 0: SwiftPM build-tool plugin to produce `mlx.metallib` directly from `libs/mlx-swift/Source/Cmlx/mlx-generated/metal/`. Today's local-dev workflow uses `scripts/fetch-metallib.sh` and CI uses the wheel-extraction step in `release-swift.yml`. +- [ ] Phase 0: SwiftPM build-tool plugin to produce `mlx.metallib` directly from `libs/mlx-swift/Source/Cmlx/mlx`. Until then, local development and CI use the canonical source builder at `scripts/fetch-metallib.sh`. - [ ] First-class `metallib_hash` field on `protocol.RegisterMessage` and `protocol.AttestationResponseMessage` (today it rides as a key inside `template_hashes`, which the coordinator stores but does not enforce). - [ ] Build-time injection of `ProviderCore.version` from a git tag (today it is a hand-bumped constant; CI consumes it as-is). diff --git a/provider-swift/Sources/ProviderBenchmark/ArrivalInvarianceBenchmark.swift b/provider-swift/Sources/ProviderBenchmark/ArrivalInvarianceBenchmark.swift index aeed5345a..d3dd50453 100644 --- a/provider-swift/Sources/ProviderBenchmark/ArrivalInvarianceBenchmark.swift +++ b/provider-swift/Sources/ProviderBenchmark/ArrivalInvarianceBenchmark.swift @@ -57,7 +57,8 @@ public struct ArrivalInvarianceBenchmarkReport: Codable, Sendable { /// launched with and the backend its engine actually built. Without it /// the phase's numbers cannot be attributed to an arm, and `.auto` /// resolves CONTIGUOUS. - public static let currentSchemaVersion = 3 + /// 4 adds required effective config-projected Gemma settings. + public static let currentSchemaVersion = 4 public let schemaVersion: Int public let modelID: String @@ -65,6 +66,8 @@ public struct ArrivalInvarianceBenchmarkReport: Codable, Sendable { public let promptTokensPerRequest: Int public let decodeTokensPerRequest: Int public let iterations: Int + /// Config-projected Gemma settings this subprocess actually benchmarked. + public let gemmaOptimizations: BenchmarkGemmaOptimizations /// Bound enforced on every measured row's `arrivalErrorMs`. Samples that /// exceed it are re-run, and the benchmark fails rather than reporting /// numbers produced by a topology it did not actually deliver. @@ -161,7 +164,8 @@ public enum ArrivalInvarianceBenchmark { iterations: Int = 3, arrivalToleranceMs: Double? = nil, maxAttemptsPerSample: Int = 3, - kvBackend: EngineV2KVBackendSelection = .auto + kvBackend: EngineV2KVBackendSelection = .auto, + gemmaOptimizations: GemmaOptimizationSettings ) async throws -> ArrivalInvarianceBenchmarkReport { let promptTokens = max(2, promptTokens) let decodeTokens = max(2, decodeTokens) @@ -203,7 +207,6 @@ public enum ArrivalInvarianceBenchmark { let engineParts = try await makeEngine( container: container, - modelDirectory: modelDirectory, isVLM: isVLM, weightBytes: facts.weightBytes, maxConcurrentRequests: patterns.map(\.delaysMs.count).max() ?? 1, @@ -321,6 +324,8 @@ public enum ArrivalInvarianceBenchmark { promptTokensPerRequest: promptTokens, decodeTokensPerRequest: decodeTokens, iterations: iterations, + gemmaOptimizations: BenchmarkGemmaOptimizations( + settings: gemmaOptimizations), arrivalToleranceMs: toleranceMs, arrivalMaxAttemptsPerSample: maxAttempts, kvBackend: BenchmarkKVBackend( @@ -504,7 +509,6 @@ public enum ArrivalInvarianceBenchmark { private static func makeEngine( container: ModelContainer, - modelDirectory: URL, isVLM: Bool, weightBytes: Int, maxConcurrentRequests: Int, @@ -525,8 +529,7 @@ public enum ArrivalInvarianceBenchmark { return try await container.perform { context -> EngineParts in let servingModel = try EngineV2Factory.benchmarkServingModel( model: context.model, - isVLM: isVLM, - modelDirectory: modelDirectory + isVLM: isVLM ) let build = try EngineV2Factory.makeProductionBuild( model: servingModel, diff --git a/provider-swift/Sources/ProviderBenchmark/BackendParityHarness.swift b/provider-swift/Sources/ProviderBenchmark/BackendParityHarness.swift index f3c293656..5b6a2c185 100644 --- a/provider-swift/Sources/ProviderBenchmark/BackendParityHarness.swift +++ b/provider-swift/Sources/ProviderBenchmark/BackendParityHarness.swift @@ -109,15 +109,9 @@ public enum BackendParityHarness { // The SERVING model is resolved EXACTLY ONCE and reused by every // engine build and by the drafter. // - // This is not an optimization. For a VLM checkpoint - // `benchmarkServingModel` runs `EngineV2VLMTextExtraction`, which - // returns a NEW model object on each call. Calling it per engine bound - // the drafter to one instance and every engine to another, and - // `CBv2MTPRoundDriver.build` could then not prove target identity — it - // returned nil and the run reported MTP inert on BOTH backends for a - // reason that was entirely the harness's. Resolving once also keeps - // the two arms on literally the same weights, so a token difference - // can only be the backend. + // Resolving once keeps the drafter and both backend arms bound to the + // exact same VLM-owned text tower, so target identity is preserved and + // a token difference can only come from the backend. struct Facts: @unchecked Sendable { let weightBytes: Int let eosTokenIds: Set @@ -134,7 +128,7 @@ public enum BackendParityHarness { let seed = ctx.tokenizer.encode( text: ThroughputSweep.seedText, addSpecialTokens: false) let servingModel = try EngineV2Factory.benchmarkServingModel( - model: ctx.model, isVLM: isVLM, modelDirectory: modelDirectory) + model: ctx.model, isVLM: isVLM) // Both halves of the packed-prefill gate are consulted by the // engine loop; only the model half is publicly readable, so read @@ -1174,7 +1168,7 @@ public enum BackendParityHarness { + "EngineV2SlotFactory). Every verdict here is a BACKEND result, not a " + "statement about how production routes a slot to that backend." + (isVLM ? " This checkpoint IS a VLM and is served here through the " - + "text-extraction seam." : "")) + + "wrapper's directly owned shared text tower." : "")) notes.append( "token comparisons are over RAW SAMPLED TOKEN IDS with temperature 0; text " + "equality is a strictly weaker oracle and is not used.") diff --git a/provider-swift/Sources/ProviderBenchmark/BenchmarkGemmaOptimizations.swift b/provider-swift/Sources/ProviderBenchmark/BenchmarkGemmaOptimizations.swift new file mode 100644 index 000000000..b4d4b7459 --- /dev/null +++ b/provider-swift/Sources/ProviderBenchmark/BenchmarkGemmaOptimizations.swift @@ -0,0 +1,19 @@ +import ProviderCore + +/// Effective config-backed Gemma posture carried by every benchmark payload. +/// +/// The low-level environment is process-local, so a parent wrapper cannot infer +/// these values from its own environment. Recording the config projection in +/// each wrapper-phase subprocess JSON makes ON/OFF artifacts attributable and +/// comparable. +public struct BenchmarkGemmaOptimizations: Codable, Sendable, Equatable { + public let prefillLayer18: Bool + public let weightedR1: Bool + public let environment: [String: String] + + public init(settings: GemmaOptimizationSettings) { + self.prefillLayer18 = settings.prefillLayer18 + self.weightedR1 = settings.weightedR1 + self.environment = GemmaOptimizationEnvironment.projection(for: settings) + } +} diff --git a/provider-swift/Sources/ProviderBenchmark/MTPProductionSession.swift b/provider-swift/Sources/ProviderBenchmark/MTPProductionSession.swift index 33e4c609c..072f0dbb5 100644 --- a/provider-swift/Sources/ProviderBenchmark/MTPProductionSession.swift +++ b/provider-swift/Sources/ProviderBenchmark/MTPProductionSession.swift @@ -5,9 +5,9 @@ import MLXVLM import ProviderCore /// Cache-only model bundle for MTP validation. It loads the target through the -/// same model factories as serving, resolves the exact VLM text model through -/// ProviderCore's production extraction seam, and loads/binds the real Gemma 4 -/// assistant. It never downloads and never contains a decoder implementation. +/// same model factories as serving, resolves the exact VLM-owned text tower, +/// and loads/binds the real Gemma 4 assistant. It never downloads and never +/// contains a decoder implementation. public final class MTPProductionModelBundle: @unchecked Sendable { public let targetID: String public let assistantID: String @@ -106,8 +106,7 @@ public final class MTPProductionModelBundle: @unchecked Sendable { convertTokenToID: { snapshot.tokenizer.convertTokenToId($0) }) let servingModel = try EngineV2Factory.benchmarkServingModel( model: snapshot.model, - isVLM: isVLM, - modelDirectory: targetDirectory) + isVLM: isVLM) guard let target = servingModel as? any Gemma4MTPTarget else { throw MTPBenchmarkError.mtpRequestedButInactive( "target model \(type(of: servingModel)) is not Gemma4MTPTarget") diff --git a/provider-swift/Sources/ProviderBenchmark/SchedulerPrefillBenchmark.swift b/provider-swift/Sources/ProviderBenchmark/SchedulerPrefillBenchmark.swift index c3f49b92c..caf4c8ef1 100644 --- a/provider-swift/Sources/ProviderBenchmark/SchedulerPrefillBenchmark.swift +++ b/provider-swift/Sources/ProviderBenchmark/SchedulerPrefillBenchmark.swift @@ -10,7 +10,8 @@ public struct SchedulerPrefillBenchmarkReport: Codable, Sendable { /// `kvBackend` block and the per-sample `resolvedKVBackend`. An /// UNVERSIONED payload predates the backend pin and cannot say which /// backend it measured, so a gate must refuse it rather than assume. - public static let currentSchemaVersion = 1 + /// 2 adds required effective config-projected Gemma settings. + public static let currentSchemaVersion = 2 public struct Sample: Codable, Sendable { public let strategy: String @@ -32,6 +33,8 @@ public struct SchedulerPrefillBenchmarkReport: Codable, Sendable { public let promptLengths: [Int] public let strategies: [String] public let iterations: Int + /// Config-projected Gemma settings this subprocess actually benchmarked. + public let gemmaOptimizations: BenchmarkGemmaOptimizations /// Selection versus the backends the measured engines were built with. public let kvBackend: BenchmarkKVBackend public let samples: [Sample] @@ -67,15 +70,16 @@ public enum SchedulerPrefillBenchmark { modelDirectory: URL, promptLengths: [Int], iterations: Int, - kvBackend: EngineV2KVBackendSelection = .auto + kvBackend: EngineV2KVBackendSelection = .auto, + gemmaOptimizations: GemmaOptimizationSettings ) async throws -> SchedulerPrefillBenchmarkReport { let lengths = promptLengths.filter { $0 > 1 }.sorted() let iterations = max(1, iterations) log("loading model \(modelID)") log(" path: \(modelDirectory.path)") - // VLM checkpoints load via the VLM factory and measure through the - // weight-sharing extracted text model (production serving path). + // VLM checkpoints load via the VLM factory and measure the exact + // text tower owned by the wrapper (the production serving path). let isVLM = ThroughputSweep.readHasVisionConfig(modelDirectory: modelDirectory) let container: ModelContainer if isVLM { @@ -109,7 +113,6 @@ public enum SchedulerPrefillBenchmark { iteration: 0, weightBytes: facts.weightBytes, isVLM: isVLM, - modelDirectory: modelDirectory, kvBackend: kvBackend ) @@ -124,7 +127,6 @@ public enum SchedulerPrefillBenchmark { iteration: iteration, weightBytes: facts.weightBytes, isVLM: isVLM, - modelDirectory: modelDirectory, kvBackend: kvBackend ) if !resolved.contains(sample.resolvedKVBackend) { @@ -143,6 +145,8 @@ public enum SchedulerPrefillBenchmark { promptLengths: lengths, strategies: [strategyLabel], iterations: iterations, + gemmaOptimizations: BenchmarkGemmaOptimizations( + settings: gemmaOptimizations), kvBackend: BenchmarkKVBackend( selection: kvBackend.rawValue, resolved: resolved), samples: samples @@ -156,7 +160,6 @@ public enum SchedulerPrefillBenchmark { iteration: Int, weightBytes: Int, isVLM: Bool, - modelDirectory: URL?, kvBackend: EngineV2KVBackendSelection ) async throws -> SchedulerPrefillBenchmarkReport.Sample { // Same KV-ceiling derivation as a single-model serving slot; far @@ -178,7 +181,7 @@ public enum SchedulerPrefillBenchmark { // to have been honoured. let parts = try await container.perform { ctx -> EngineParts in let servingModel = try EngineV2Factory.benchmarkServingModel( - model: ctx.model, isVLM: isVLM, modelDirectory: modelDirectory) + model: ctx.model, isVLM: isVLM) let build = try EngineV2Factory.makeProductionBuild( model: servingModel, tokenizer: ctx.tokenizer, diff --git a/provider-swift/Sources/ProviderBenchmark/ThroughputSweep.swift b/provider-swift/Sources/ProviderBenchmark/ThroughputSweep.swift index cd7386e47..ec4acf410 100644 --- a/provider-swift/Sources/ProviderBenchmark/ThroughputSweep.swift +++ b/provider-swift/Sources/ProviderBenchmark/ThroughputSweep.swift @@ -31,6 +31,12 @@ public enum ThroughputSweep { public static let defaultDecodePromptTokens = 64 public static let defaultDecodeIterations = 1 + /// Throughput cells must generate the requested budget for every row. + /// Honoring model EOS would compare different token counts and lets one + /// early-stopping row corrupt a batch aggregate; arrival invariance uses + /// the same fixed-budget contract. + static let fixedBudgetStopTokens: Set = [] + /// Snapshot of model facts read once, off-actor, inside `perform`. private struct ModelFacts: Sendable { let weightBytes: Int @@ -64,15 +70,15 @@ public enum ThroughputSweep { decodePromptTokens: Int = defaultDecodePromptTokens, decodeIterations: Int = defaultDecodeIterations, kvBackend: EngineV2KVBackendSelection = .auto, + gemmaOptimizations: GemmaOptimizationSettings, hardware: HardwareInfo, efficiency: Double = DecodeBandwidthModel.defaultBandwidthEfficiency ) async throws -> ThroughputSweepReport { log("loading model \(modelID)") log(" path: \(modelDirectory.path)") - // VLM checkpoints (config declares `vision_config`) load through the - // VLM factory and serve through the weight-sharing extracted text - // model — the same construction every production slot performs. + // VLM checkpoints load through the VLM factory and serve through the + // exact text tower owned by that wrapper, matching production. let isVLM = readHasVisionConfig(modelDirectory: modelDirectory) let container: ModelContainer if isVLM { @@ -111,7 +117,6 @@ public enum ThroughputSweep { iterations: decodeIterations, weightBytes: facts.weightBytes, isVLM: isVLM, - modelDirectory: modelDirectory, kvBackend: kvBackend ) let decode = decodeOutcome.samples @@ -159,6 +164,8 @@ public enum ThroughputSweep { decode: decode, derived: derived, notes: notes, + gemmaOptimizations: BenchmarkGemmaOptimizations( + settings: gemmaOptimizations), kvBackend: ThroughputSweepReport.KVBackend( selection: kvBackend.rawValue, resolved: decodeOutcome.resolvedBackends), @@ -321,7 +328,6 @@ public enum ThroughputSweep { iterations: Int, weightBytes: Int, isVLM: Bool, - modelDirectory: URL?, kvBackend: EngineV2KVBackendSelection ) async -> DecodeOutcome { let sizes = batchSizes.filter { $0 > 0 }.sorted() @@ -339,7 +345,7 @@ public enum ThroughputSweep { let warmUp = await runDecodeBatch( container: container, modelID: modelID, baseTokens: baseTokens, batchSize: 1, decodeTokens: 4, promptLen: promptLen, weightBytes: weightBytes, - isVLM: isVLM, modelDirectory: modelDirectory, kvBackend: kvBackend) + isVLM: isVLM, kvBackend: kvBackend) // The warm-up is the FIRST cell to hit a refused paged selection, so // it carries the reason even when the sized cells below fail // identically. Keep it: an operator should not have to infer the @@ -351,8 +357,7 @@ public enum ThroughputSweep { let (totalTokens, maxElapsed, resolved, failure, submitFailure) = await runDecodeBatch( container: container, modelID: modelID, baseTokens: baseTokens, batchSize: batchSize, decodeTokens: genTokens, promptLen: promptLen, - weightBytes: weightBytes, isVLM: isVLM, modelDirectory: modelDirectory, - kvBackend: kvBackend) + weightBytes: weightBytes, isVLM: isVLM, kvBackend: kvBackend) if outcome.record(resolved), let resolved { log(" engine resolved kv backend: \(resolved)") } @@ -413,7 +418,6 @@ public enum ThroughputSweep { promptLen: Int, weightBytes: Int, isVLM: Bool, - modelDirectory: URL?, kvBackend: EngineV2KVBackendSelection ) async -> ( totalTokens: Int, maxElapsed: Duration, resolvedBackend: String?, @@ -430,17 +434,16 @@ public enum ThroughputSweep { UInt64(Int.max))) struct EngineParts: @unchecked Sendable { let engine: any CBv2Engine - let eosTokenIds: Set /// The backend the factory resolved to, with any fallback reason. let resolvedBackend: String } let parts: EngineParts do { parts = try await container.perform { ctx -> EngineParts in - // Serving-model resolution: VLM checkpoints run the - // weight-sharing text extraction, exactly like a slot build. + // Serving-model resolution matches production: VLM checkpoints + // use the exact text tower owned by the loaded wrapper. let servingModel = try EngineV2Factory.benchmarkServingModel( - model: ctx.model, isVLM: isVLM, modelDirectory: modelDirectory) + model: ctx.model, isVLM: isVLM) // `makeProductionBuild` is the construction // `makeProductionEngine` wraps, and additionally hands back the // backend kind the engine actually resolved to — the fact a @@ -453,7 +456,6 @@ public enum ThroughputSweep { kvBackend: kvBackend) return EngineParts( engine: build.engine, - eosTokenIds: ctx.configuration.eosTokenIds, resolvedBackend: build.resolvedKVBackendDescriptor) } } catch { @@ -465,7 +467,6 @@ public enum ThroughputSweep { return (0, .zero, nil, "\(error)", nil) } let engine = parts.engine - let eosTokenIds = parts.eosTokenIds let result = await withTaskGroup(of: RowMeasure.self) { group -> (Int, Duration, String?) in @@ -482,7 +483,7 @@ public enum ThroughputSweep { promptTokens: prompt, sampling: CBv2SamplingParams(temperature: 0.0), maxTokens: decodeTokens + 1, - stopTokens: eosTokenIds + stopTokens: Self.fixedBudgetStopTokens )) } catch { Self.log(" submit failed: \(error)") @@ -585,8 +586,8 @@ public enum ThroughputSweep { return (sorted[middle - 1] + sorted[middle]) / 2 } - /// Whether the checkpoint's config.json declares a `vision_config` - /// (VLM — load via the VLM factory, serve via the text extraction). + /// Whether config.json declares `vision_config` (load through VLMModelFactory + /// and serve through the wrapper-owned text tower). static func readHasVisionConfig(modelDirectory: URL) -> Bool { let url = modelDirectory.appendingPathComponent("config.json") guard let data = try? Data(contentsOf: url), diff --git a/provider-swift/Sources/ProviderBenchmark/ThroughputSweepReport.swift b/provider-swift/Sources/ProviderBenchmark/ThroughputSweepReport.swift index 36e9931cd..7ddf8700e 100644 --- a/provider-swift/Sources/ProviderBenchmark/ThroughputSweepReport.swift +++ b/provider-swift/Sources/ProviderBenchmark/ThroughputSweepReport.swift @@ -1,4 +1,5 @@ import Foundation +import ProviderCore /// Machine-readable result of `darkbloom benchmark --sweep`. /// @@ -15,7 +16,8 @@ public struct ThroughputSweepReport: Codable, Sendable { /// `decode[].resolvedKVBackend`. /// 4 adds the required `decodeCoverage` block: which cells were ASKED /// for versus which ones actually produced a measurement. - public static let currentSchemaVersion = 4 + /// 5 adds required effective config-projected Gemma settings. + public static let currentSchemaVersion = 5 public struct Hardware: Codable, Sendable { public let chipName: String @@ -229,6 +231,8 @@ public struct ThroughputSweepReport: Codable, Sendable { public let decode: [DecodeSample] public let derived: Derived public let notes: [String] + /// Config-projected Gemma settings this subprocess actually benchmarked. + public let gemmaOptimizations: BenchmarkGemmaOptimizations /// Selection versus resolved backend. Always present since schema 3: a /// decode curve whose backend is unknown is not comparable to anything. public let kvBackend: KVBackend @@ -251,6 +255,7 @@ public struct ThroughputSweepReport: Codable, Sendable { decode: [DecodeSample], derived: Derived, notes: [String], + gemmaOptimizations: BenchmarkGemmaOptimizations, kvBackend: KVBackend = KVBackend(selection: "auto", resolved: []), decodeConstructionFailure: DecodeConstructionFailure? = nil, decodeCoverage: DecodeCoverage = DecodeCoverage( @@ -264,6 +269,7 @@ public struct ThroughputSweepReport: Codable, Sendable { self.decode = decode self.derived = derived self.notes = notes + self.gemmaOptimizations = gemmaOptimizations self.kvBackend = kvBackend self.decodeConstructionFailure = decodeConstructionFailure self.decodeCoverage = decodeCoverage diff --git a/provider-swift/Sources/ProviderCore/Config/BetaFeatures.swift b/provider-swift/Sources/ProviderCore/Config/BetaFeatures.swift index 6f0aeb812..e494f813b 100644 --- a/provider-swift/Sources/ProviderCore/Config/BetaFeatures.swift +++ b/provider-swift/Sources/ProviderCore/Config/BetaFeatures.swift @@ -1,6 +1,6 @@ import Foundation -/// A user-facing opt-in *beta* feature. +/// A user-facing configurable *beta* feature. /// /// Beta features are intentionally **config-backed**, not environment-variable /// backed: the launchd daemon started by `darkbloom start` only inherits a tiny @@ -22,6 +22,13 @@ public struct BetaFeature: Sendable, Identifiable { public let details: String /// Whether `darkbloom restart` is required for a change to take effect. public let requiresRestart: Bool + /// Where the backing field lives in `provider.toml` + /// (`[section]` + `key =`), so the CLI can tell "already at the requested + /// value AND pinned in the file" apart from "same value via decode + /// default". An absent key must be WRITTEN on an explicit toggle: the + /// operator asked for the value durably, not for today's decode default + /// (a future default flip would otherwise silently move their provider). + public let configAddress: (section: String, key: String)? private let read: @Sendable (ProviderConfig) -> Bool private let write: @Sendable (Bool, inout ProviderConfig) -> Void @@ -32,6 +39,7 @@ public struct BetaFeature: Sendable, Identifiable { summary: String, details: String, requiresRestart: Bool, + configAddress: (section: String, key: String)? = nil, read: @escaping @Sendable (ProviderConfig) -> Bool, write: @escaping @Sendable (Bool, inout ProviderConfig) -> Void ) { @@ -40,6 +48,7 @@ public struct BetaFeature: Sendable, Identifiable { self.summary = summary self.details = details self.requiresRestart = requiresRestart + self.configAddress = configAddress self.read = read self.write = write } @@ -55,13 +64,51 @@ public struct BetaFeature: Sendable, Identifiable { } } -/// The registry of opt-in beta features. +/// The registry of configurable beta features. /// /// Adding a beta toggle = adding one ``BetaFeature`` entry here (and its backing -/// `ProviderConfig` field). The `darkbloom beta` command and `darkbloom status` -/// are driven entirely off this list, so they need no per-feature code. +/// `ProviderConfig` field). Features declare their own defaults; the +/// `darkbloom beta` command and `darkbloom status` are driven entirely off this +/// list, so they need no per-feature code. public enum BetaFeatures { public static let all: [BetaFeature] = [ + BetaFeature( + id: "gemma-prefill-layer18", + title: "Gemma layer-18 prefill submission", + summary: "Default ON. Submit prefill work every 18 layers; disable for legacy submission behavior.", + details: """ + Default ON for existing and new configs. Writes prefill_layer18 \ + under [gemma_optimizations]; provider config is authoritative over \ + the low-level process environment. Restart after changing it. \ + Disable and restart to restore the legacy one-final-submission \ + prefill behavior. + """, + requiresRestart: true, + configAddress: (section: "gemma_optimizations", key: "prefill_layer18"), + read: { $0.gemmaOptimizations.prefillLayer18 }, + write: { enabled, config in + config.gemmaOptimizations.prefillLayer18 = enabled + } + ), + BetaFeature( + id: "gemma-weighted-r1", + title: "Gemma weighted unsort + safe R1", + summary: "Default ON. Coupled weighted-unsort and safe-R1 expert paths with one rollback.", + details: """ + Default ON for existing and new configs. Writes weighted_r1 under \ + [gemma_optimizations]. This single production control keeps direct \ + weighted expert reduction coupled to the safe exact-shape R1 QMM \ + path; neither half can be selected independently. Provider config \ + is authoritative over the low-level process environment. Disable \ + and restart to restore both legacy paths together. + """, + requiresRestart: true, + configAddress: (section: "gemma_optimizations", key: "weighted_r1"), + read: { $0.gemmaOptimizations.weightedR1 }, + write: { enabled, config in + config.gemmaOptimizations.weightedR1 = enabled + } + ), BetaFeature( id: "mtp", title: "Multi-token prediction (speculative decoding)", @@ -75,6 +122,7 @@ public enum BetaFeatures { mtp_drafter_path under [backend] to a local drafter directory. """, requiresRestart: true, + configAddress: (section: "backend", key: "mtp"), read: { $0.backend.mtp }, write: { enabled, config in config.backend.mtp = enabled } ), diff --git a/provider-swift/Sources/ProviderCore/Config/GemmaOptimizationEnvironment.swift b/provider-swift/Sources/ProviderCore/Config/GemmaOptimizationEnvironment.swift new file mode 100644 index 000000000..6d2ec02bd --- /dev/null +++ b/provider-swift/Sources/ProviderCore/Config/GemmaOptimizationEnvironment.swift @@ -0,0 +1,89 @@ +import Darwin + +/// Projects the config-backed Gemma controls into the low-level environment +/// consumed while MLX and MLX-LM initialize process-wide optimization state. +public enum GemmaOptimizationEnvironment { + public static let prefillLayer18Key = "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL" + public static let weightedUnsortKey = "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT" + public static let safeR1Key = "MLX_GATHER_QMM_EXPERT_SLICES" + + /// Return the complete production projection for `settings`. + /// + /// Weighted unsort and safe R1 always receive the same value: production + /// never exposes either half of the benchmark-selected pair independently. + public static func projection( + for settings: GemmaOptimizationSettings + ) -> [String: String] { + let weightedR1 = settings.weightedR1 ? "1" : "0" + return [ + prefillLayer18Key: settings.prefillLayer18 ? "18" : "0", + weightedUnsortKey: weightedR1, + safeR1Key: weightedR1, + ] + } + + /// Raised when the process refuses one or more projected controls. + /// + /// Weighted unsort and safe R1 are process-start latches, so a projection + /// that only partially applied leaves the coupled pair in a state no + /// benchmark ever measured. Callers must abort startup instead of + /// continuing on the surviving half. + public struct ApplicationFailure: Error, Equatable, CustomStringConvertible { + /// Rejected keys, sorted, so the message is stable across runs. + public let keys: [String] + /// `errno` reported by the first rejected key. + public let code: Int32 + + public init(keys: [String], code: Int32) { + self.keys = keys + self.code = code + } + + public var description: String { + let reason = strerror(code).map { String(cString: $0) } + ?? "errno \(code)" + return """ + failed to apply Gemma optimization controls \ + [\(keys.joined(separator: ", "))]: \(reason) + """ + } + } + + /// Apply the complete projection to the current process. + /// + /// Config is authoritative, so every key is overwritten even when the + /// launching shell supplied a conflicting low-level value. Throws + /// `ApplicationFailure` if the environment rejects any key. + public static func apply(_ settings: GemmaOptimizationSettings) throws { + try apply(settings) { name, value, overwrite in + errno = 0 + guard setenv(name, value, overwrite) == 0 else { + return errno == 0 ? EINVAL : errno + } + return 0 + } + } + + /// Apply every projected key, then report all rejections at once. + /// + /// - Parameter set: applies one key and returns `0` on success or the + /// failing `errno` otherwise. + static func apply( + _ settings: GemmaOptimizationSettings, + set: (_ name: String, _ value: String, _ overwrite: Int32) -> Int32 + ) throws { + var rejected: [String] = [] + var code: Int32 = 0 + // Sorted so a failure reports the same keys regardless of the + // dictionary's per-process hash ordering. + for (name, value) in projection(for: settings).sorted(by: { $0.key < $1.key }) { + let status = set(name, value, 1) + guard status != 0 else { continue } + rejected.append(name) + if code == 0 { code = status } + } + guard rejected.isEmpty else { + throw ApplicationFailure(keys: rejected, code: code) + } + } +} diff --git a/provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift b/provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift new file mode 100644 index 000000000..52957b8c7 --- /dev/null +++ b/provider-swift/Sources/ProviderCore/Config/GemmaOptimizationSettings.swift @@ -0,0 +1,36 @@ +/// Production controls for the benchmark-selected Gemma optimization stack. +/// +/// Both controls default on so provider configs written before these keys were +/// introduced receive the selected stack. Operators can roll either part back +/// in `provider.toml`; the change takes effect after the provider restarts. +public struct GemmaOptimizationSettings: Sendable, Equatable, Codable { + /// Submit prefill work every 18 Gemma transformer layers. + public var prefillLayer18: Bool + + /// Enable the coupled weighted-unsort and safe-R1 expert paths. + /// + /// These paths intentionally share one production control. Exposing them + /// independently could select a combination that was not benchmarked. + public var weightedR1: Bool + + public init( + prefillLayer18: Bool = true, + weightedR1: Bool = true + ) { + self.prefillLayer18 = prefillLayer18 + self.weightedR1 = weightedR1 + } + + enum CodingKeys: String, CodingKey { + case prefillLayer18 = "prefill_layer18" + case weightedR1 = "weighted_r1" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.prefillLayer18 = try container.decodeIfPresent( + Bool.self, forKey: .prefillLayer18) ?? true + self.weightedR1 = try container.decodeIfPresent( + Bool.self, forKey: .weightedR1) ?? true + } +} diff --git a/provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift b/provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift index 305bb091c..7f5c9eb2a 100644 --- a/provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift +++ b/provider-swift/Sources/ProviderCore/Config/ProviderConfig.swift @@ -10,6 +10,7 @@ /// - Backend settings (port, model, continuous batching, idle timeout) /// - Coordinator connection settings (URL, heartbeat interval) /// - Scheduling windows +/// - Config-backed Gemma optimization controls /// /// A default config is generated based on detected hardware when the provider /// is first initialized. CLI flags can override config values at runtime. @@ -332,6 +333,7 @@ public struct ProviderConfig: Sendable, Equatable, Codable { public var backend: BackendSettings public var coordinator: CoordinatorSettings public var schedule: ScheduleConfig? + public var gemmaOptimizations: GemmaOptimizationSettings /// Schema version of the `provider.toml` this config came from /// (`config_version`, top level, written by the startup stamp in /// `migrateConfigIfNeeded`). @@ -366,12 +368,14 @@ public struct ProviderConfig: Sendable, Equatable, Codable { backend: BackendSettings = BackendSettings(), coordinator: CoordinatorSettings = CoordinatorSettings(), schedule: ScheduleConfig? = nil, + gemmaOptimizations: GemmaOptimizationSettings = GemmaOptimizationSettings(), configVersion: Int = ProviderConfig.currentConfigVersion ) { self.provider = provider self.backend = backend self.coordinator = coordinator self.schedule = schedule + self.gemmaOptimizations = gemmaOptimizations self.configVersion = configVersion } @@ -380,6 +384,7 @@ public struct ProviderConfig: Sendable, Equatable, Codable { case backend case coordinator case schedule + case gemmaOptimizations = "gemma_optimizations" case configVersion = "config_version" } @@ -389,6 +394,9 @@ public struct ProviderConfig: Sendable, Equatable, Codable { var backend = try container.decodeIfPresent(BackendSettings.self, forKey: .backend) ?? BackendSettings() self.coordinator = try container.decodeIfPresent(CoordinatorSettings.self, forKey: .coordinator) ?? CoordinatorSettings() self.schedule = try container.decodeIfPresent(ScheduleConfig.self, forKey: .schedule) + self.gemmaOptimizations = try container.decodeIfPresent( + GemmaOptimizationSettings.self, forKey: .gemmaOptimizations + ) ?? GemmaOptimizationSettings() // One-time concurrency-default migrations, selected by the stamp the // file carries. @@ -513,6 +521,13 @@ public enum ConfigManager: Sendable { } /// Load config from a file path. + /// + /// This is the production file-loading boundary: a file that EXISTS but + /// cannot decode throws `ConfigError.parseFailed` (see + /// ``parseValidating(_:)``) instead of silently falling back to defaults. + /// Callers that want missing-file leniency check `fileExists` first; only + /// a NONEXISTENT path defaults (`readFailed` is thrown here, and the + /// snapshot/`loadDefault` layers substitute defaults in that case). public static func load(from path: URL) throws -> ProviderConfig { let content: String do { @@ -520,7 +535,7 @@ public enum ConfigManager: Sendable { } catch { throw ConfigError.readFailed(path: path.path, underlying: error) } - return parse(content) + return try parseValidating(content) } /// Load config from the default path. Returns default config if file doesn't exist. @@ -568,6 +583,13 @@ public enum ConfigManager: Sendable { // MARK: - TOML parsing /// Parse a TOML string into a ProviderConfig. + /// + /// LENIENT, test-facing entry point: ANY decode failure falls back to a + /// whole-config default, exactly matching historical behavior. Production + /// file loads must NOT use this — a malformed `[gemma_optimizations]` + /// entry (e.g. `weighted_r1 = 0` as an integer) would otherwise silently + /// re-enable the whole default-on optimization stack with zero log. Use + /// ``parseValidating(_:)`` (via ``load(from:)``) on that path. public static func parse(_ content: String) -> ProviderConfig { do { return try TOMLDecoder().decode(ProviderConfig.self, from: content) @@ -581,6 +603,22 @@ public enum ConfigManager: Sendable { } } + /// Parse a TOML string into a ProviderConfig, failing loudly. + /// + /// Unlike ``parse(_:)``, any decode failure throws + /// `ConfigError.parseFailed` carrying the decoder's description, so an + /// operator who fat-fingers `provider.toml` is told at startup instead of + /// unknowingly serving on whole-config defaults. Missing / partial content + /// still decodes with per-key defaults (default-on Gemma stack) — only an + /// undecodable FILE is rejected. + public static func parseValidating(_ content: String) throws -> ProviderConfig { + do { + return try TOMLDecoder().decode(ProviderConfig.self, from: content) + } catch { + throw ConfigError.parseFailed(detail: "\(error)") + } + } + /// Serialize a ProviderConfig to the provider's TOML config format. public static func serialize(_ config: ProviderConfig) -> String { do { diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2Config.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2Config.swift index d40baab46..c33f30fcb 100644 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2Config.swift +++ b/provider-swift/Sources/ProviderCore/Inference/EngineV2Config.swift @@ -83,9 +83,6 @@ public enum EngineV2RefusalReason: String, Sendable { /// (`EngineV2ProductionError.unsupportedModel`) — should be unreachable /// behind the scan-time supported-set gate; kept as loud insurance. case unsupportedModel = "unsupported_model" - /// The Gemma 4 VLM text-model extraction failed (config decode, weight - /// re-key, verify, or the forward-parity gate). - case vlmExtractionFailed = "vlm_extraction_failed" /// A load-time KV re-slice would push some co-resident slot below the /// minimum serviceable grant (`EngineV2KVSizing` floor). case resliceFloor = "reslice_floor" @@ -121,8 +118,6 @@ public enum EngineV2RefusalReason: String, Sendable { return .pagedBackendUnavailable case EngineV2ProductionError.invalidPagedPoolDType: return .pagedKVDTypeInvalid - case is EngineV2VLMTextExtractionError: - return .vlmExtractionFailed default: return .engineInitFailed } diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Benchmark.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Benchmark.swift index 3f29226cb..4c80e0fc8 100644 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Benchmark.swift +++ b/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Benchmark.swift @@ -1,33 +1,20 @@ // Copyright © 2026 Eigen Labs. // -// Benchmark-facing seam: the perf-gate harness (`ProviderBenchmark`'s -// ThroughputSweep / SchedulerPrefillBenchmark) must measure the SERVING -// model — for Gemma 4 VLM checkpoints that is the weight-sharing -// CBv2-adapted text model produced by `EngineV2VLMTextExtraction`, exactly -// as `EngineV2SlotFactory.makeProductionBridge` builds it. Without this the -// sweep would hand the raw VLM wrapper to `makeProductionEngine` and refuse -// (`unsupportedModel`), measuring nothing. +// Benchmark-facing seam: perf harnesses must measure the exact module served +// in production. Gemma 4 VLM checkpoints expose their directly owned +// `Gemma4TextModel`; no extraction, re-keying, or second module is involved. -import Foundation import MLXLMCommon extension EngineV2Factory { /// Resolve the CBv2-serving model for a loaded checkpoint: the model - /// itself for text checkpoints, the weight-sharing extracted text model - /// for VLM checkpoints (zero extra weight memory; load-time parity gate - /// included — throws on any extraction/verify failure). + /// itself for text checkpoints, or the exact VLM-owned text tower for + /// Gemma 4 VLM checkpoints. public static func benchmarkServingModel( model: any LanguageModel, - isVLM: Bool, - modelDirectory: URL? + isVLM: Bool ) throws -> any LanguageModel { - guard isVLM else { return model } - guard let modelDirectory else { - throw EngineV2VLMTextExtractionError.missingModelDirectory - } - return try EngineV2VLMTextExtraction.extractTextModel( - from: model, modelDirectory: modelDirectory - ).model + try directServingModel(model: model, isVLM: isVLM) } } diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Production.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Production.swift index 8e81190be..1c1a83547 100644 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Production.swift +++ b/provider-swift/Sources/ProviderCore/Inference/EngineV2Factory+Production.swift @@ -39,14 +39,14 @@ import Foundation import MLX import MLXLLM import MLXLMCommon +import MLXVLM /// Failure modes of production v2-engine construction. Each maps to the /// factory's REFUSAL path (ERROR `engine_v2_refusal` telemetry + throw). enum EngineV2ProductionError: Error, CustomStringConvertible { /// The loaded module is not a CBv2-adapted family (an unexpected - /// architecture). Allowlisted Gemma 4 VLM wrappers do NOT land here — - /// the slot factory extracts their CBv2-adapted text model first - /// (`EngineV2VLMTextExtraction`) and hands THAT to this factory. + /// architecture). Gemma 4 VLM wrappers are resolved to their directly + /// owned text tower before engine construction. case unsupportedModel(String) /// No KV byte budget is left under the unified-memory cap — an engine /// admitted with a zero ceiling would reject every request, so the @@ -66,10 +66,11 @@ enum EngineV2ProductionError: Error, CustomStringConvertible { /// (`EngineV2Factory.pagedPoolDType(environment:)`) and surfaced as a /// REFUSAL only for an EXPLICIT `.paged` selection — the measurement /// posture the knob exists for, where silently serving fp16 under an - /// fp32 label would fake a control arm. Under `.auto` (the fleet - /// default since v0.8.0) the factory catches it and DEGRADES to - /// contiguous with `fallbackReason = "invalid_dtype: …"` instead: one - /// typo'd env var must not 503 every slot on the fleet. + /// fp32 label would fake a control arm. If `.auto` resolves paged, the + /// factory catches this and DEGRADES to contiguous with + /// `fallbackReason = "invalid_dtype: …"` instead. That path is dormant + /// while `.auto` resolves contiguous as of v0.8.1, but remains the safety + /// contract for any future paged default. case invalidPagedPoolDType(String) var description: String { @@ -89,6 +90,20 @@ enum EngineV2ProductionError: Error, CustomStringConvertible { } extension EngineV2Factory { + /// Resolve the exact module instance served by CBv2. Gemma 4 VLM owns + /// its `Gemma4TextModel`; direct VLM forwards and CBv2 therefore share + /// one language tower, one parameter tree, and one residency footprint. + static func directServingModel( + model: any LanguageModel, isVLM: Bool + ) throws -> any LanguageModel { + guard isVLM else { return model } + guard let gemma4 = model as? MLXVLM.Gemma4 else { + throw EngineV2ProductionError.unsupportedModel( + String(describing: type(of: model))) + } + return gemma4.textModel + } + /// Clamp a KV admission ceiling to physical unified memory. A ceiling /// above physical RAM can only come from a mis-derivation upstream; the @@ -283,8 +298,8 @@ extension EngineV2Factory { /// Build the real `EngineV2` over a loaded model. /// /// - Parameters: - /// - model: the loaded language module (the SAME instance the legacy - /// engine serves — weights are shared, never duplicated). + /// - model: the loaded serving language module; for Gemma 4 VLM this + /// is the exact text tower owned by the wrapper. /// - tokenizer: the model's tokenizer, for incremental detokenization. /// - kvBytesCapacity: admission ceiling for live sequence KV, in bytes /// (derive from `UnifiedMemoryCap.kvBudgetBytes`). @@ -532,11 +547,11 @@ extension EngineV2Factory { let layerKinds: [CBv2LayerKind] let newCaches: ((Int, CBv2LayerKind) -> any CBv2AttendingLayerCache) - -> [any CBv2AttendingLayerCache] + throws -> [any CBv2AttendingLayerCache] switch model { case let gemma as Gemma4TextModel: layerKinds = gemma.cbv2LayerKinds - newCaches = { make in gemma.newCacheV2(makeLayerCache: make) } + newCaches = { make in try gemma.newCacheV2(makeLayerCache: make) } case let gptoss as GPTOSSModel: layerKinds = gptoss.cbv2LayerKinds newCaches = { make in gptoss.newCacheV2(makeLayerCache: make) } @@ -749,10 +764,10 @@ extension EngineV2Factory { let schedulerConfig = CBv2SchedulerConfig( maxConcurrentRequests: max(1, maxConcurrentRequests)) - func contiguousPreparation() -> ProductionBackendPreparation { + func contiguousPreparation() throws -> ProductionBackendPreparation { let backend = CBv2ContiguousKVBackend( config: CBv2ContiguousBackendConfig(bytesCapacity: cappedCapacity)) - let caches = newCaches { index, kind in + let caches = try newCaches { index, kind in CBv2LayerCache(layerIndex: index, kind: kind) } return ProductionBackendPreparation( @@ -833,7 +848,7 @@ extension EngineV2Factory { // exists, so no admitted page is ever unbacked. slabCommitment: plan.commitment) let pagedCaches = paged.makeLayerCaches() - let caches = newCaches { index, _ in pagedCaches[index] } + let caches = try newCaches { index, _ in pagedCaches[index] } return ProductionBackendPreparation( model: model, maxConcurrentRequests: maxConcurrentRequests, @@ -866,7 +881,7 @@ extension EngineV2Factory { } } } - return contiguousPreparation() + return try contiguousPreparation() } /// Emergency rollback kill-switch for the monotonic deadline leases. Set diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory+MTP.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory+MTP.swift index ecd3a7484..e5c7d1ee0 100644 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory+MTP.swift +++ b/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory+MTP.swift @@ -1,4 +1,3 @@ -import Foundation import MLXLMCommon /// Model handle + EOS config snapshot pulled out of `ModelContainer.perform`. @@ -9,8 +8,8 @@ struct EngineV2ModelSnapshot: @unchecked Sendable { let extraEOSTokens: [String] } -/// Target extraction plus fail-open assistant preparation, completed before KV -/// re-slicing so final sizing uses retained assistant bytes. +/// Direct serving-target resolution plus fail-open assistant preparation, +/// completed before KV re-slicing so sizing uses retained assistant bytes. struct EngineV2PreparedModel: @unchecked Sendable { let snapshot: EngineV2ModelSnapshot let servingModel: any LanguageModel @@ -45,45 +44,31 @@ extension EngineV2SlotFactory { private static func servingModel( modelId: String, isVLM: Bool, - modelDirectory: URL?, snapshot: EngineV2ModelSnapshot, emitTelemetry: (@Sendable (TelemetryEvent) -> Void)?, logInfo: @escaping @Sendable (String) -> Void ) throws -> any LanguageModel { guard isVLM else { return snapshot.model } - guard let modelDirectory else { - let error = EngineV2VLMTextExtractionError.missingModelDirectory - EngineV2Factory.emitRefusalTelemetry( - modelId: modelId, - reason: .vlmExtractionFailed, - error: error, - emitTelemetry: emitTelemetry) - throw error - } - let extraction: EngineV2VLMTextExtraction.Extraction do { - extraction = try EngineV2VLMTextExtraction.extractTextModel( - from: snapshot.model, modelDirectory: modelDirectory) + let target = try EngineV2Factory.directServingModel( + model: snapshot.model, isVLM: true) + logInfo( + "engine_v2: \(modelId) using the Gemma 4 VLM-owned text tower " + + "directly (shared identity and residency)") + return target } catch { EngineV2Factory.emitRefusalTelemetry( modelId: modelId, - reason: .vlmExtractionFailed, + reason: EngineV2RefusalReason.classify(error), error: error, emitTelemetry: emitTelemetry) throw error } - if let parityDiff = extraction.parityMaxAbsLogitDiff { - logInfo( - "engine_v2: \(modelId) VLM text-model extraction passed the " - + "load-time forward parity gate (max |Δlogit| \(parityDiff))") - } - return extraction.model } static func prepareProductionModel( modelId: String, isVLM: Bool, - modelDirectory: URL?, container: ModelContainer, specDecPreparation: SpecDecPreparation, assistantLoader: any ProviderMTPAssistantLoading = Gemma4ProviderMTPAssistantLoader(), @@ -95,7 +80,6 @@ extension EngineV2SlotFactory { let servingModel = try servingModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, snapshot: snapshot, emitTelemetry: emitTelemetry, logInfo: logInfo) @@ -153,7 +137,6 @@ extension EngineV2SlotFactory { static func prepareRecoveryModel( modelId: String, isVLM: Bool, - modelDirectory: URL?, container: ModelContainer, previousArtifact: SpecDecArtifact?, previousStatus: MTPActivationStatus, @@ -192,7 +175,6 @@ extension EngineV2SlotFactory { let target = try servingModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, snapshot: snapshot, emitTelemetry: emitTelemetry, logInfo: logInfo) @@ -208,7 +190,6 @@ extension EngineV2SlotFactory { let target = try servingModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, snapshot: snapshot, emitTelemetry: emitTelemetry, logInfo: logInfo) diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory.swift index be35d1dad..27ea067bc 100644 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory.swift +++ b/provider-swift/Sources/ProviderCore/Inference/EngineV2SlotFactory.swift @@ -7,12 +7,11 @@ // slots through THIS one path so the assembly can never drift between // them: snapshot the loaded module's EOS config out of the container, // apply the model-specific EOS policy (`ModelEOSPolicy`), build the -// production CBv2 engine over the loaded module (with the weight-sharing -// VLM text extraction for Gemma 4 VLM checkpoints), and wrap it in an -// `EngineV2Bridge` via the fail-loud `EngineV2Factory.makeBridge` (any -// construction failure emits the ERROR `engine_v2_refusal` telemetry and -// throws — the caller unloads and maps to a 503; there is no legacy -// fallback). +// production CBv2 engine over the loaded module (using the Gemma 4 VLM's +// directly owned text tower), and wrap it in an `EngineV2Bridge` via the +// fail-loud `EngineV2Factory.makeBridge` (any construction failure emits the +// ERROR `engine_v2_refusal` telemetry and throws — the caller unloads and +// maps to a 503; there is no legacy fallback). // // Call-site differences stay at the call sites: the ProviderLoop // registers the bridge with `EngineV2Runtime` (heartbeat/cancel fan-out) @@ -21,10 +20,124 @@ // are limited to `makeEngineOverride` (scripted engines, no weights). import Foundation +import MLX import MLXLLM import MLXLMCommon import ProviderCoreFoundation +enum GemmaOptimizationReason: String, Sendable, Equatable { + case disabled + case modelIneligible = "model_ineligible" + case aotUnavailable = "aot_unavailable" + case naxPrecedence = "nax_precedence" + case effective +} + +struct GemmaOptimizationState: Sendable, Equatable { + let name: String + let requested: Bool + let effective: Bool + let reason: GemmaOptimizationReason + + var compactDescription: String { + "\(name)(requested=\(requested),effective=\(effective),reason=\(reason.rawValue))" + } +} + +/// Pure requested/effective resolution for the three retained Gemma controls. +/// Safe R1 is inferred from one unarmed device snapshot; this type never +/// resets, arms, or samples route counters. +struct GemmaOptimizationReport: Sendable, Equatable { + let layer18: GemmaOptimizationState + let weightedUnsort: GemmaOptimizationState + let safeR1: GemmaOptimizationState + + init( + layer18Requested: Bool, + layer18Effective: Bool, + weightedUnsortRequested: Bool, + weightedUnsortEffective: Bool, + safeR1Requested: Bool, + safeR1GeometryEligible: Bool, + safeR1AOTAvailable: Bool, + safeR1NAXAvailable: Bool + ) { + layer18 = Self.resolve( + name: "layer18", + requested: layer18Requested, + modelEligible: layer18Effective) + weightedUnsort = Self.resolve( + name: "weighted_unsort", + requested: weightedUnsortRequested, + modelEligible: weightedUnsortEffective) + safeR1 = Self.resolve( + name: "safe_r1", + requested: safeR1Requested, + modelEligible: safeR1GeometryEligible, + aotAvailable: safeR1AOTAvailable, + naxAvailable: safeR1NAXAvailable) + } + + var states: [GemmaOptimizationState] { + [layer18, weightedUnsort, safeR1] + } + + func logLine(modelId: String) -> String { + "engine_v2: \(modelId) gemma optimizations " + + states.map(\.compactDescription).joined(separator: " ") + } + + func telemetryEvents(modelId: String) -> [TelemetryEvent] { + states.map { state in + var event = TelemetryEvent( + source: .provider, + severity: .info, + kind: .engineHealth, + message: "engine_v2: gemma optimization " + + state.compactDescription) + event.fields = TelemetryFieldFilter.filter([ + "component": .string("engine"), + "operation": .string("gemma_optimization_\(state.name)"), + "backend": .string("engine_v2"), + "model": .string(modelId), + // Existing allowlisted field, carrying the bounded 2-bit + // requested/effective state without a telemetry schema change. + "target": .string( + "requested_\(state.requested ? 1 : 0)_effective_" + + "\(state.effective ? 1 : 0)"), + "reason": .string(state.reason.rawValue), + ]) + return event + } + } + + private static func resolve( + name: String, + requested: Bool, + modelEligible: Bool, + aotAvailable: Bool? = nil, + naxAvailable: Bool = false + ) -> GemmaOptimizationState { + let reason: GemmaOptimizationReason + if !requested { + reason = .disabled + } else if !modelEligible { + reason = .modelIneligible + } else if aotAvailable == false { + reason = .aotUnavailable + } else if naxAvailable { + reason = .naxPrecedence + } else { + reason = .effective + } + return GemmaOptimizationState( + name: name, + requested: requested, + effective: reason == .effective, + reason: reason) + } +} + enum EngineV2SlotFactory { /// Narrow assembly seams for production-order regression tests. Normal @@ -62,10 +175,9 @@ enum EngineV2SlotFactory { /// - Parameters: /// - modelId: catalog id the slot serves under. /// - modelType: `model_type` from config.json (EOS policy input). - /// - isVLM: config declares `vision_config` — the engine is built - /// over `EngineV2VLMTextExtraction`'s weight-sharing text model. - /// - modelDirectory: checkpoint dir (required for VLM extraction). - /// - container: the just-loaded model container. + /// - isVLM: config declares `vision_config` — the engine directly uses + /// the `Gemma4TextModel` owned by the loaded VLM wrapper. + /// - modelDirectory: checkpoint dir (prompt-contract identity input). /// - tokenizer: the container's tokenizer handle. /// - sizing: scheduler-free sizing snapshot (fp16 KV rate, context, /// default max tokens). @@ -86,7 +198,7 @@ enum EngineV2SlotFactory { /// `ProviderLoop.EngineV2SlotHooks`); nil ⇒ the real /// `EngineV2Factory.makeProductionEngine`. SSD cache instances and /// stats logging exist only on the production path. - /// - logInfo: sink for the VLM parity-gate + cache-state info lines. + /// - logInfo: sink for shared-tower + cache-state info lines. /// - logWarning: sink for the both-tiers-requested WARN line. static func makeProductionBridge( modelId: String, @@ -134,8 +246,7 @@ enum EngineV2SlotFactory { } /// Production bundle assembly. Assistant preparation is deliberately - /// fail-open; target extraction/engine construction retain their existing - /// fail-loud semantics. + /// fail-open; direct target resolution and engine construction fail loud. static func makeProductionBundle( modelId: String, modelType: String?, @@ -205,7 +316,6 @@ enum EngineV2SlotFactory { prepared = try await prepareProductionModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, container: container, specDecPreparation: SpecDecPreparation( artifact: nil, status: specDecPreparation.status), @@ -217,7 +327,6 @@ enum EngineV2SlotFactory { prepared = try await prepareProductionModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, container: container, specDecPreparation: specDecPreparation, assistantLoader: assistantLoader, @@ -288,12 +397,9 @@ enum EngineV2SlotFactory { // only RAM claims are per-request staging reservations in the // shared `GlobalKVCacheBudget` (refused ⇒ silent recompute). // - // VLM slots derive layer kinds from config.json's text_config - // alone (`EngineV2VLMTextExtraction.cbv2LayerKinds` — drift tests - // pin config-derived shape == engine truth); the weight-sharing - // extraction itself still runs inside engine construction. A - // family with no derivable kinds gets no cache (it would throw - // unsupportedModel at engine build anyway). + // VLM slots use the layer kinds of the exact text tower already + // resolved from the loaded wrapper. A family with no adapted serving + // model gets no prepared backend and is refused before cache creation. var ssdPrefixCache: SSDPrefixCache? var cacheCapability: CBv2PrefixReuseCapability? var cacheConstructionStatus = PrefixCacheConstructionStatus.configDisabled @@ -489,6 +595,31 @@ enum EngineV2SlotFactory { await bridge.startSSDPrefixCacheStatsLogger(cache: ssdPrefixCache) } await bridge.configureMTPStatus(mtpStatus) + if let gemmaModel = servingModel as? Gemma4TextModel { + // One load-time snapshot only. Never arm the benchmark counters in + // production: the QMM hot path remains free of counter atomics. + let r1 = GPU.gemma4ExpertQMMDiagnostics() + let layerInterval = gemmaModel.cbv2PrefillChunkEvalInterval + let report = GemmaOptimizationReport( + layer18Requested: layerInterval > 0, + layer18Effective: + layerInterval > 0 + && gemmaModel.cbv2LayerKinds.count >= layerInterval, + weightedUnsortRequested: gemmaModel.weightedExpertUnsortRequested, + weightedUnsortEffective: gemmaModel.weightedExpertUnsortEffective, + safeR1Requested: r1.requested, + safeR1GeometryEligible: gemmaModel.expertQMMGeometryEligible, + safeR1AOTAvailable: r1.aotAvailable, + safeR1NAXAvailable: r1.naxAvailable) + logInfo(report.logLine(modelId: modelId)) + for event in report.telemetryEvents(modelId: modelId) { + if let emitTelemetry { + emitTelemetry(event) + } else { + TelemetryClient.shared.emit(event) + } + } + } logInfo( "engine_v2: \(modelId) prefix cache " + prefixCacheStateDescription( diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2SupportedModels.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2SupportedModels.swift index 2daf10966..91e28fd28 100644 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2SupportedModels.swift +++ b/provider-swift/Sources/ProviderCore/Inference/EngineV2SupportedModels.swift @@ -10,10 +10,9 @@ // declares (the value `ModelScanner.parseModelInfo` stamps on `ModelInfo`): // // * `gpt_oss` — GPT-OSS (GPTOSSModel) -// * `gemma4` — Gemma 4 VLM wrapper +// * `gemma4` — Gemma 4 VLM wrapper, serving through its directly +// owned text tower plus vision prefill // * `gemma4_text` — Gemma 4 text target -// wrapper (`gemma4`, served via the weight-sharing -// text-model extraction + vision prefill) // // Everything else (gemma3, qwen*, llama, …) has no CBv2 adapter: it is // dropped from the advertised set at startup and at prefetch-verify time diff --git a/provider-swift/Sources/ProviderCore/Inference/EngineV2VLMTextExtraction.swift b/provider-swift/Sources/ProviderCore/Inference/EngineV2VLMTextExtraction.swift deleted file mode 100644 index 16a661f0f..000000000 --- a/provider-swift/Sources/ProviderCore/Inference/EngineV2VLMTextExtraction.swift +++ /dev/null @@ -1,369 +0,0 @@ -// Copyright © 2026 Eigen Labs. -// -// ContinuousBatchingV2 — weight-sharing text-model extraction for VLM slots. -// -// Every production Gemma 4 checkpoint ships a vision tower, so the provider -// loads it through `VLMModelFactory` and the resident module is MLXVLM's -// `Gemma4` wrapper — whose language model is a PRIVATE inline duplicate of -// the text architecture ("MLXVLM can't import MLXLLM") with none of the -// CBv2 hooks (`cbv2LayerKinds` / `newCacheV2`). The v2 engine therefore -// could never serve a VLM slot directly, and before v0.7.2 the per-slot -// `isVLM` gate excluded 100% of prod Gemma traffic from engine v2. -// -// This file builds the CBv2-adapted MLXLLM `Gemma4TextModel` OVER THE SAME -// WEIGHT ARRAYS the wrapper already holds: -// -// 1. decode the checkpoint's `config.json` `text_config` with MLXLLM's -// `Gemma4TextConfiguration` decoder and construct a lazy skeleton -// (nothing is materialized — MLXArray init is lazy until eval); -// 2. re-apply the checkpoint's quantization STRUCTURE to the skeleton the -// exact way `loadWeights` did for the wrapper (scales-presence gate + -// the same per-layer table from `BaseConfiguration`, whose keys live in -// the checkpoint's `language_model.`-prefixed key space); -// 3. re-key the wrapper's live parameter tree (`language_model.model.X` → -// `model.X`, `language_model.lm_head.X` → `lm_head.X`), run it through -// `Gemma4TextModel.sanitize` (drops shared-KV k/v duplicates the -// wrapper allocates but the MLXLLM model does not), and -// `update(parameters:verify:[.all])` — missing/extra/mis-shaped keys -// all THROW, which the engine factory catches as the standard -// `engine_v2_refusal` ERROR + load failure (never silent wrongness); -// 4. run a tiny forward through BOTH the wrapper's text path and the -// extracted model and require cross-containment of each side's greedy -// argmax in the other side's top-5, plus a bounded max |Δlogit| — a -// load-time backstop against catastrophic extraction bugs (see -// `assertForwardParity` for why bit-parity is structurally -// unattainable; env `DARKBLOOM_ENGINE_V2_VLM_PARITY_CHECK=0` skips). -// -// The result is a SEPARATE module instance (own norms/rope/layer objects) -// sharing only the immutable parameter arrays — zero extra weight memory, -// and no module-level mutable state shared with the wrapper, which is what -// makes the Qwen3.5-mrope class of cross-path state corruption structurally -// impossible here. Concurrency: the legacy vision forward (wrapper) and v2 -// text forward (extracted model) may interleave on the same arrays; both -// are read-only over the parameters (forward passes never mutate weights) -// and each path owns its private KV caches. - -import Foundation -import MLX -import MLXLLM -import MLXLMCommon -import MLXNN -import MLXVLM - -/// Failure modes of VLM text-model extraction. Every case lands in -/// `EngineV2Factory.makeBridge`'s catch → ERROR `engine_v2_refusal` -/// telemetry + legacy serving. The messages are operator-facing (they ride -/// the telemetry `error` field), so they say exactly what to look at. -enum EngineV2VLMTextExtractionError: Error, CustomStringConvertible { - /// The loaded module is not a VLM wrapper this extraction understands. - case unsupportedWrapper(String) - /// The slot factory could not hand us the model directory (needed for - /// `config.json`). - case missingModelDirectory - /// `config.json` was unreadable or had no decodable `text_config`. - case invalidConfig(String) - /// The checkpoint has quantized weights but no `quantization` block in - /// `config.json` to derive the skeleton's quantization structure from. - case missingQuantizationConfig - /// The load-time forward parity gate failed: the extracted text model - /// disagrees with the wrapper's own text path on the same weights. - case parityMismatch(String) - - var description: String { - switch self { - case .unsupportedWrapper(let type): - return "engine_v2 vlm extraction: unsupported VLM wrapper \(type)" - case .missingModelDirectory: - return "engine_v2 vlm extraction: model directory unavailable for config.json" - case .invalidConfig(let detail): - return "engine_v2 vlm extraction: config.json unusable (\(detail))" - case .missingQuantizationConfig: - return "engine_v2 vlm extraction: quantized weights but no quantization block in config.json" - case .parityMismatch(let detail): - return "engine_v2 vlm extraction: wrapper/extracted forward parity failed (\(detail))" - } - } -} - -/// Weight-sharing extraction of the CBv2-adapted MLXLLM text model from a -/// loaded MLXVLM wrapper. Pure functions; no state. -enum EngineV2VLMTextExtraction { - - /// Env kill switch for the load-time forward parity gate ("0"/"false"/ - /// "no"/"off" disables). Default ON — the check is one tiny prefill at - /// model-load time and is the backstop against silent architecture - /// drift between MLXVLM's inline text model and MLXLLM's. - static let parityCheckFlag = "DARKBLOOM_ENGINE_V2_VLM_PARITY_CHECK" - - /// Checkpoint key-space prefix of the wrapper's language model. Both the - /// parameter tree and the per-layer quantization table use it. - private static let languageModelPrefix = "language_model." - - /// Result of one extraction: the CBv2-adapted text model (sharing the - /// wrapper's weight arrays) plus the parity probe's max |Δlogit| for the - /// slot factory's log line (nil when the parity gate was disabled). - struct Extraction { - let model: Gemma4TextModel - let parityMaxAbsLogitDiff: Float? - } - - /// Build an MLXLLM `Gemma4TextModel` over the weight arrays of a loaded - /// MLXVLM `Gemma4` wrapper. See the file header for the full mechanism. - /// - /// - Parameters: - /// - model: the slot's loaded module (must be `MLXVLM.Gemma4`). - /// - modelDirectory: the checkpoint directory (for `config.json`). - /// - environment: env snapshot (parity-gate kill switch). - static func extractTextModel( - from model: any LanguageModel, - modelDirectory: URL, - environment: [String: String] = ProcessInfo.processInfo.environment - ) throws -> Extraction { - guard let wrapper = model as? MLXVLM.Gemma4 else { - throw EngineV2VLMTextExtractionError.unsupportedWrapper( - String(describing: type(of: model))) - } - - // 1. Text config: decode the checkpoint's `text_config` with the SAME - // decoder `LLMModelFactory` would use for a text-only checkpoint. - let configURL = modelDirectory.appendingPathComponent("config.json") - let configData: Data - do { - configData = try Data(contentsOf: configURL) - } catch { - throw EngineV2VLMTextExtractionError.invalidConfig( - "read \(configURL.path): \(error)") - } - let textConfig = try decodeTextConfiguration(configData: configData) - - // Quantization table: `BaseConfiguration` holds the checkpoint-wide - // default plus the per-layer overrides, keyed in the checkpoint's - // `language_model.`-prefixed key space. - let baseConfig = try? JSONDecoder.json5().decode(BaseConfiguration.self, from: configData) - - // 2-3. Lazy skeleton → quantization structure → weight-sharing update. - let skeleton = Gemma4TextModel(textConfig) - let textWeights = reKeyedTextWeights(wrapper: wrapper, sanitizer: skeleton) - try applyQuantizationStructure( - skeleton: skeleton, weights: textWeights, - perLayerQuantization: baseConfig?.perLayerQuantization) - // verify: [.all] — a missing model key, an unused weight key, or a - // shape mismatch all throw here. That is the design: any drift - // between the wrapper's parameter tree and the MLXLLM architecture - // must fail LOUDLY at load (→ engine_v2_refusal ERROR + 503), - // never produce a silently wrong serving model. - try skeleton.update( - parameters: ModuleParameters.unflattened(textWeights), verify: [.all]) - - // 4. Load-time forward parity gate (env-gated, default on). - var parityDiff: Float? = nil - if parityCheckEnabled(environment: environment) { - // Return the probe's transient buffers to the OS before the load - // path's post-build admission/headroom reads. The probe's two - // forwards leave ~GiBs of intermediates in the MLX pool; fence - // async GPU completion first (M4 IOKit guard), and clean up on - // the parity-failure throw path too. - defer { - MLX.Stream().synchronize() - MLX.Memory.clearCache() - } - parityDiff = try assertForwardParity( - wrapper: wrapper, extracted: skeleton, vocabSize: textConfig.vocabSize) - } - - return Extraction(model: skeleton, parityMaxAbsLogitDiff: parityDiff) - } - - /// A VLM checkpoint's CBv2 layer kinds from `config.json`'s - /// `text_config` ALONE — identical to what the extracted text model - /// reports (the drift tests pin config-derived shape == engine truth). - /// Used by the slot factory to construct the SSD prefix cache for VLM - /// slots (layout-epoch + adoption-bound binding) BEFORE the extraction - /// runs inside engine construction. nil when the config is - /// unreadable/undecodable — no cache is built (the extraction will - /// throw on the same config moments later). - static func cbv2LayerKinds(modelDirectory: URL) -> [CBv2LayerKind]? { - let configURL = modelDirectory.appendingPathComponent("config.json") - guard let configData = try? Data(contentsOf: configURL), - let textConfig = try? decodeTextConfiguration(configData: configData) - else { return nil } - return textConfig.cbv2LayerKinds - } - - // MARK: - Steps - - /// Decode MLXLLM's `Gemma4TextConfiguration` from the VLM checkpoint's - /// `text_config` block. The top-level `quantization` block is merged in - /// so the configuration's informational `quantizationBits`/`GroupSize` - /// fields reflect the checkpoint (they do not drive the skeleton's - /// quantization — `applyQuantizationStructure` does). - static func decodeTextConfiguration(configData: Data) throws -> Gemma4TextConfiguration { - let root: [String: Any] - do { - guard - let parsed = try JSONSerialization.jsonObject(with: configData) - as? [String: Any] - else { - throw EngineV2VLMTextExtractionError.invalidConfig("config.json is not an object") - } - root = parsed - } catch let error as EngineV2VLMTextExtractionError { - throw error - } catch { - throw EngineV2VLMTextExtractionError.invalidConfig("parse config.json: \(error)") - } - guard var textConfigJSON = root["text_config"] as? [String: Any] else { - throw EngineV2VLMTextExtractionError.invalidConfig("no text_config object") - } - if textConfigJSON["quantization"] == nil, let quantization = root["quantization"] { - textConfigJSON["quantization"] = quantization - } - do { - let textConfigData = try JSONSerialization.data(withJSONObject: textConfigJSON) - return try JSONDecoder.json5().decode(Gemma4TextConfiguration.self, from: textConfigData) - } catch { - throw EngineV2VLMTextExtractionError.invalidConfig("decode text_config: \(error)") - } - } - - /// Re-key the wrapper's live parameter tree into the text model's key - /// space and drop everything outside the language model (vision tower, - /// multimodal embedder). The result then passes through the text - /// model's own `sanitize`, which drops the k/v-projection duplicates - /// the wrapper allocates for KV-shared layers (MLXLLM does not allocate - /// those modules) and leaves everything else untouched. - private static func reKeyedTextWeights( - wrapper: MLXVLM.Gemma4, sanitizer: Gemma4TextModel - ) -> [String: MLXArray] { - var textWeights: [String: MLXArray] = [:] - for (key, value) in wrapper.parameters().flattened() { - guard key.hasPrefix(languageModelPrefix) else { continue } - textWeights[String(key.dropFirst(languageModelPrefix.count))] = value - } - return sanitizer.sanitize(weights: textWeights) - } - - /// Mirror `loadWeights`' quantization pass onto the skeleton: a module is - /// quantized iff the (re-keyed, live) weights carry `.scales`, with - /// (groupSize, bits, mode) resolved from the checkpoint's per-layer table - /// under the module's `language_model.`-prefixed checkpoint key — the - /// exact lookup the wrapper's own load performed, so the two module trees - /// can never disagree on quantization structure. All arrays involved stay - /// lazy; `update(parameters:)` replaces them before anything evaluates. - private static func applyQuantizationStructure( - skeleton: Gemma4TextModel, - weights: [String: MLXArray], - perLayerQuantization: BaseConfiguration.PerLayerQuantization? - ) throws { - let hasQuantizedWeights = weights.keys.contains { $0.hasSuffix(".scales") } - guard hasQuantizedWeights else { return } - guard let perLayerQuantization else { - throw EngineV2VLMTextExtractionError.missingQuantizationConfig - } - quantize(model: skeleton) { path, _ in - guard weights["\(path).scales"] != nil else { return nil } - return perLayerQuantization - .quantization(layer: languageModelPrefix + path)?.asTuple - } - } - - // MARK: - Parity gate - - static func parityCheckEnabled(environment: [String: String]) -> Bool { - guard - let raw = environment[parityCheckFlag]? - .trimmingCharacters(in: .whitespaces).lowercased(), !raw.isEmpty - else { return true } - return !["0", "false", "no", "off"].contains(raw) - } - - /// Top-k window for the bidirectional argmax-containment check. - static let parityTopK = 5 - /// Ceiling on max |Δlogit| across the probe. Final logits are - /// softcapped to ±`final_logit_softcapping` (30 on every Gemma 4 - /// checkpoint), so unrelated distributions differ by up to ~60; - /// same-weights implementation noise measured ≤ ~8 (see below). - static let parityMaxAbsLogitDiff: Float = 20 - - /// One tiny forward through the wrapper's text path and the extracted - /// model on identical tokens — a CATASTROPHIC-EXTRACTION detector, not - /// a bit-parity check. - /// - /// Token-exact parity between the two is structurally unattainable - /// (measured on gemma-4-26B-A4B qat-4bit, 2026-07): - /// - /// * the two implementations have different bf16 kernel/fusion - /// orderings, giving ~0.5 max |Δlogit| even at position 0 (where - /// RoPE is the identity), which flips near-tie top-8-of-128 MoE - /// expert selections at later positions (~8 max |Δlogit|); - /// * MLXVLM's wrapper mis-implements the checkpoint's declared - /// `rope_type: "proportional"` for full-attention layers as an - /// HF-style truncated-dims partial RoPE (freqs /128, partner +64), - /// while MLXLLM's `ProportionalRoPE` implements the declared type - /// (freqs /512 over the full head, partner +256 — verified against - /// the mlx_lm Python reference). The extracted model keeps the - /// CORRECT scheme; the wrapper discrepancy is flagged upstream. - /// - /// What a correct extraction guarantees — and what this gate enforces — - /// is that both models compute the same function up to implementation - /// noise: each side's greedy argmax must sit inside the other side's - /// top-`parityTopK` at EVERY position, and max |Δlogit| must stay under - /// `parityMaxAbsLogitDiff`. A mis-extracted model (wrong keys, wrong - /// quantization structure, wrong config) produces unrelated - /// distributions and fails both with overwhelming probability. The - /// probe is fixed and both models are deterministic, so for a given - /// (checkpoint, binary) the gate either always passes or always fails — - /// no per-load flakiness. - private static func assertForwardParity( - wrapper: MLXVLM.Gemma4, extracted: Gemma4TextModel, vocabSize: Int - ) throws -> Float { - // Fixed probe: small ids well inside every Gemma vocab; length stays - // far under the sliding window so the check exercises both layer - // types without materializing meaningful KV. - let probeTokens = [2, 651, 6134, 1024, 578, 108, 2364].map { - min($0, max(0, vocabSize - 1)) - } - let inputs = MLXArray(probeTokens.map(Int32.init)).expandedDimensions(axis: 0) - - let wrapperLogits = wrapper(inputs, cache: nil).asType(.float32) - let extractedLogits = extracted(inputs, cache: nil).asType(.float32) - - let k = parityTopK - // Top-k token ids per position, [1, L, k] (unordered within k). - let wrapperTopK = argPartition(-wrapperLogits, kth: k - 1, axis: -1)[.ellipsis, .. ProviderConfig { + try TOMLDecoder().decode( + ProviderConfig.self, from: retainedConfigurationTOML) + } + + static func validateRetainedProjection( + _ projection: [String: String] + ) throws { + let expected = [ + GemmaOptimizationEnvironment.prefillLayer18Key: "18", + GemmaOptimizationEnvironment.weightedUnsortKey: "1", + GemmaOptimizationEnvironment.safeR1Key: "1", + ] + guard projection == expected else { + throw VerificationError.unexpectedProjection(projection) + } + } + + static func rejectedEnvironmentKeys( + projection: [String: String], + environment: [String: String] + ) -> [String] { + projection.compactMap { key, value in + environment[key] == value ? nil : key + }.sorted() + } + + static func validateSafeR1( + requested: Bool, + aotAvailable: Bool, + countersArmed: Bool + ) throws { + guard requested else { throw VerificationError.safeR1NotRequested } + guard aotAvailable else { throw VerificationError.safeR1AOTUnavailable } + guard !countersArmed else { throw VerificationError.safeR1CountersArmed } + } + + static func containsGemmaOptimizationSuccessMarker(_ output: Data) -> Bool { + guard let text = String(data: output, encoding: .utf8) else { return false } + return text.split(whereSeparator: \.isNewline).contains { + $0 == Substring(gemmaOptimizationSuccessMarker) + } + } public static func runPagedKernel( shapes: [PagedAttentionKernelSmokeShape] = diff --git a/provider-swift/Sources/ProviderCore/Inference/SlotSizingSnapshot.swift b/provider-swift/Sources/ProviderCore/Inference/SlotSizingSnapshot.swift index 41fc582d9..1d3602a68 100644 --- a/provider-swift/Sources/ProviderCore/Inference/SlotSizingSnapshot.swift +++ b/provider-swift/Sources/ProviderCore/Inference/SlotSizingSnapshot.swift @@ -116,7 +116,6 @@ public struct SlotSizingSnapshot: Sendable, Equatable { struct ModuleFacts: @unchecked Sendable { let bytes: Int let moduleKVRate: Int? - let isGemma4VLMWrapper: Bool } let facts = await container.perform { ctx -> ModuleFacts in let bytes = ctx.model.parameters().flattened().reduce(0) { $0 + $1.1.nbytes } @@ -124,27 +123,24 @@ public struct SlotSizingSnapshot: Sendable, Equatable { // layer kinds (GPT-OSS derives them from the LOADED trunk, so // they are congruent with the actual layers even when config.json // omits `layer_types`). - var rate: Int? = nil - var isWrapper = false + let rate: Int? switch ctx.model { case let gemma as Gemma4TextModel: rate = fp16KVBytesPerToken(layerKinds: gemma.cbv2LayerKinds) case let gptoss as GPTOSSModel: rate = fp16KVBytesPerToken(layerKinds: gptoss.cbv2LayerKinds) - case is MLXVLM.Gemma4: - // The VLM wrapper has no CBv2 hooks; its extracted text model - // does. Derive from the checkpoint's text_config below (the - // SAME decoder the extraction uses, so the kinds equal what - // the extracted model will report). - isWrapper = true + case let gemma as MLXVLM.Gemma4: + // Direct ownership makes the loaded tower engine truth with + // no config re-decode or second topology that can drift. + rate = fp16KVBytesPerToken(layerKinds: gemma.textModel.cbv2LayerKinds) default: - break + rate = nil } - return ModuleFacts(bytes: bytes, moduleKVRate: rate, isGemma4VLMWrapper: isWrapper) + return ModuleFacts(bytes: bytes, moduleKVRate: rate) } - // Architecture metadata (context window + the config-parse fallback - // rate) from config.json — the same parse the load-gate estimate uses. + // Architecture metadata (context window + non-CBv2 fallback rate) + // comes from the same config parse the load-gate estimate uses. let architecture: ModelArchitecture if let modelPath { architecture = KVEstimation.parseModelArchitecture( @@ -154,15 +150,10 @@ public struct SlotSizingSnapshot: Sendable, Equatable { } var kvRate = facts.moduleKVRate ?? 0 - if kvRate <= 0, facts.isGemma4VLMWrapper, let modelPath { - kvRate = gemma4VLMTextKVRate(modelDirectory: modelPath) ?? 0 - } if kvRate <= 0 { - // Non-CBv2 module (or a wrapper whose text_config failed to - // decode): fall back to the config-parse figure so callers that - // only need a rough rate (vision reservations in tests) still - // get one. Such a model cannot build a v2 engine and is refused - // at load — this value never sizes a real engine grant. + // Non-CBv2 module: fall back to the config-parse figure so + // callers that only need a rough rate still get one. Such a + // model cannot build a v2 engine and is refused at load. kvRate = KVEstimation.resolvedKVBytesPerToken( architecture: architecture, weightBytes: facts.bytes) } @@ -216,16 +207,4 @@ public struct SlotSizingSnapshot: Sendable, Equatable { return total } - /// CBv2 layer kinds for a Gemma 4 VLM checkpoint's TEXT model, decoded - /// from `config.json`'s `text_config` with the same decoder the - /// weight-sharing extraction uses (`EngineV2VLMTextExtraction`), so the - /// kinds equal what the extracted `Gemma4TextModel` will report. - static func gemma4VLMTextKVRate(modelDirectory: URL) -> Int? { - let configURL = modelDirectory.appendingPathComponent("config.json") - guard let configData = try? Data(contentsOf: configURL), - let textConfig = try? EngineV2VLMTextExtraction.decodeTextConfiguration( - configData: configData) - else { return nil } - return fp16KVBytesPerToken(layerKinds: textConfig.cbv2LayerKinds) - } } diff --git a/provider-swift/Sources/ProviderCore/Process/BoundedProcess.swift b/provider-swift/Sources/ProviderCore/Process/BoundedProcess.swift index c2be2f736..5322f15cd 100644 --- a/provider-swift/Sources/ProviderCore/Process/BoundedProcess.swift +++ b/provider-swift/Sources/ProviderCore/Process/BoundedProcess.swift @@ -33,6 +33,55 @@ enum BoundedProcess { environment: [String: String]? = nil, timeout: TimeInterval, captureStderrTail: Int = 0 + ) throws { + try runProcess( + executable, + arguments: arguments, + environment: environment, + standardOutput: FileHandle.nullDevice, + timeout: timeout, + captureStderrTail: captureStderrTail) + } + + /// Run a bounded child while retaining stdout for signed-artifact + /// assertions. A temporary file avoids the pipe backpressure deadlock that + /// would otherwise let a verbose child fill its pipe before termination. + static func runCapturingStandardOutput( + _ executable: URL, + arguments: [String], + environment: [String: String]? = nil, + timeout: TimeInterval + ) throws -> Data { + let outputURL = FileManager.default.temporaryDirectory + .appendingPathComponent("darkbloom-child-\(UUID().uuidString).stdout") + guard FileManager.default.createFile( + atPath: outputURL.path, contents: nil) + else { + throw CocoaError(.fileWriteUnknown) + } + defer { try? FileManager.default.removeItem(at: outputURL) } + + let output = try FileHandle(forWritingTo: outputURL) + defer { try? output.close() } + try runProcess( + executable, + arguments: arguments, + environment: environment, + standardOutput: output, + timeout: timeout, + captureStderrTail: 0) + try output.synchronize() + try output.close() + return try Data(contentsOf: outputURL) + } + + private static func runProcess( + _ executable: URL, + arguments: [String], + environment: [String: String]?, + standardOutput: Any, + timeout: TimeInterval, + captureStderrTail: Int ) throws { let process = Process() process.executableURL = executable @@ -42,7 +91,7 @@ enum BoundedProcess { environment, uniquingKeysWith: { _, override in override }) } - process.standardOutput = FileHandle.nullDevice + process.standardOutput = standardOutput // A discarded stderr costs more than it saves. When the child fails, // its own message is the entire actionable content -- a preflight diff --git a/provider-swift/Sources/ProviderCore/ProviderCore.swift b/provider-swift/Sources/ProviderCore/ProviderCore.swift index cf096f474..1f1ae0ce1 100644 --- a/provider-swift/Sources/ProviderCore/ProviderCore.swift +++ b/provider-swift/Sources/ProviderCore/ProviderCore.swift @@ -75,26 +75,13 @@ public enum ProviderCore { // jinja_null_bridge / jinja_template / model_load), so durable telemetry can // tell the two indistinguishable gpt-oss 500 modes apart. Wire-compatible: // `error_reason` is an optional inference-error field, omitted when nil. - // 0.7.2 lets engine_v2 (continuous batching) serve TEXT requests on - // allowlisted VLM-loaded Gemma 4 slots. Every prod Gemma 4 checkpoint - // ships a vision tower, so it loads via VLMModelFactory and the per-slot - // isVLM gate previously kept 100% of Gemma traffic on the legacy engine. - // The slot factory now extracts the CBv2-adapted MLXLLM text model over - // the SAME weight arrays (zero extra weight memory) and serves text - // through v2; image/video requests keep the legacy VLM path. No protocol - // change — capability is behavioral, gated by the existing engine_v2 - // allowlist + flag. - // 0.7.3 fixes the v0.7.2 black-hole incident: the VLM text extraction's - // two module trees each lazily built their own multi-GiB SwitchGLU fused - // gate+up expert cache at the load-time parity probe (~15 GiB × 2 on - // gemma-4-26b-8bit), pushing 64 GB (8-bit) and 36 GB (qat-4bit) boxes - // past the 90% unified-memory cap so the shared KV gate rejected every - // request forever. The trees now share ONE fused cache - // (SwitchGLU.shareFusedGateUpCache); the load path re-checks serveable - // KV headroom AFTER the engine build and unloads instead of advertising - // a dead model; and GlobalKVCacheBudget audits + drops stale - // reservations under sustained full-rejection (defense in depth). No - // protocol changes. + // 0.7.2 first enabled CBv2 for VLM-loaded Gemma 4 by reconstructing a + // separate MLXLLM text module over shared arrays. That convention was the + // source of the historical multi-tree memory/parity complexity. + // Gemma 4 VLM now owns the canonical `Gemma4TextModel` directly: direct + // VLM, CBv2, media prefill, and MTP all retain the same module identity. + // The post-build serveable-headroom guard and stale-reservation cleanup + // remain defense in depth; no protocol change is required. // 0.7.5 is the ONE-ENGINE release: every request — text, image, video, // mixed — on every slot serves through ContinuousBatchingV2, and the // legacy BatchScheduler engine is DELETED from the binary (~15k lines @@ -210,7 +197,15 @@ public enum ProviderCore { // staged, matched or adopted where adoption diverges. Paged code, the // DARKBLOOM_CBV2_PAGED_KV kill switch, the crash-loop guard and the // blocking paged CI lane all stay; `engine_v2_kv_backend = "paged"` - // still resolves paged. `engine_v2_max_concurrent = 8` is NOT coupled - // to paged and stays at 8. - public static let version = "0.8.1" + // still resolves paged. The box-wide concurrency default returns to 4 + // with contiguous. The v0.8.1 migration preserves an existing 8 only for + // an explicitly paged config, and 8 remains the supported upper bound for + // operator and per-model overrides. + // + // 0.8.2 adds the benchmark-retained Gemma 4 optimization stack: layer-18 + // lazy prefill submission plus the coupled weighted-unsort/safe-R1 path, + // both default-on and durably rollbackable through provider config. VLM, + // CBv2, media prefill, and MTP share the canonical Gemma text tower, and + // serve/benchmark startup projects config before the first MLX access. + public static let version = "0.8.2" } diff --git a/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2.swift b/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2.swift index fedc313ec..85e134ffc 100644 --- a/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2.swift +++ b/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2.swift @@ -224,11 +224,11 @@ extension ProviderLoop { /// Re-slice KV grants for the newcomer + existing slots, shrink /// existing engines, and build the newcomer's bridge. On ANY throw — - /// re-slice floor, extraction failure, engine construction — the - /// newcomer's weights are RELEASED (box + clearCache) and only THEN is - /// every existing engine's grant RESTORED exactly (Codex-review - /// ordering: restoring first would let Σ(grants) exceed the true fleet - /// budget while the failed newcomer's weights are still resident). + /// direct serving-model resolution, re-slice floor, or engine construction + /// — the newcomer's weights are RELEASED (box + clearCache) and only THEN + /// every existing engine's grant is RESTORED exactly (restoring first + /// would let Σ(grants) exceed the true fleet budget while the failed + /// newcomer's weights are still resident). /// Returns the bridge, already registered with `engineV2Runtime`. /// CALLER HOLDS the re-slice gate (see `acquireResliceGate`), spanning /// through slot installation, so a concurrent idle-timeout unload's @@ -259,9 +259,9 @@ extension ProviderLoop { return build.bundle.bridge } - /// MTP-aware load funnel. Target extraction and assistant load/bind complete - /// before final sizing and re-slicing, so fallback removes the prospective - /// assistant charge and active MTP adds it exactly once. + /// MTP-aware load funnel. Direct target resolution and assistant load/bind + /// complete before final sizing and re-slicing, so fallback removes the + /// prospective assistant charge and active MTP adds it exactly once. internal func resliceAndBuildEngineV2Bundle( modelId: String, modelType: String?, @@ -279,7 +279,6 @@ extension ProviderLoop { prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, container: newcomerBox.borrow(), specDecPreparation: specDecPreparation, assistantLoader: engineV2SlotHooks?.assistantLoader @@ -487,12 +486,11 @@ extension ProviderLoop { /// caller unloads and maps to 503. There is no legacy path. /// /// VLM slots: every production Gemma 4 checkpoint ships a vision tower, - /// so the loaded module is MLXVLM's wrapper — which has no CBv2 hooks. - /// The engine is built over `EngineV2VLMTextExtraction`'s weight-sharing - /// MLXLLM text model: TEXT requests serve through v2, and ALL media — - /// image, video, mixed — prefills through v2 via `EngineV2VisionPrefill` - /// (v0.7.5; media construction failures REFUSE loudly, 503 — see - /// `MultiModelBatchSchedulerEngine.streamChatCompletion`). + /// so the loaded MLXVLM wrapper directly exposes its owned MLXLLM text + /// tower to CBv2. Text, image, video, mixed, and MTP requests therefore + /// share one language-model identity; media prefills use + /// `EngineV2VisionPrefill` (construction failures refuse loudly with 503; + /// see `MultiModelBatchSchedulerEngine.streamChatCompletion`). /// /// On success the bridge is registered with `engineV2Runtime` BEFORE the /// caller installs the slot, so a request routed the instant the slot diff --git a/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2Liveness.swift b/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2Liveness.swift index aa4934392..e41e18bac 100644 --- a/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2Liveness.swift +++ b/provider-swift/Sources/ProviderCore/ProviderLoop+EngineV2Liveness.swift @@ -187,7 +187,6 @@ extension ProviderLoop { var prepared = try await EngineV2SlotFactory.prepareRecoveryModel( modelId: modelId, isVLM: slot.isVLM, - modelDirectory: modelDirectory, container: slot.container, previousArtifact: slot.engineBundle.mtpArtifact, previousStatus: slot.engineBundle.mtpStatus, diff --git a/provider-swift/Sources/ProviderCore/ProviderLoop+ModelLoading.swift b/provider-swift/Sources/ProviderCore/ProviderLoop+ModelLoading.swift index be23fc589..570c4d07f 100644 --- a/provider-swift/Sources/ProviderCore/ProviderLoop+ModelLoading.swift +++ b/provider-swift/Sources/ProviderCore/ProviderLoop+ModelLoading.swift @@ -555,13 +555,13 @@ extension ProviderLoop { var sizing = slotBuild.sizing var engineV2Bridge = engineBundle.bridge - // Post-BRIDGE measured-headroom re-guard (v0.7.3, kept): the - // engine build (and, for VLM slots, the text-model extraction + - // parity probe) can retain additional load-time memory beyond - // the weights the check above measured. Re-measure so a box - // whose full load-time footprint leaves no serveable KV unloads - // and 503s instead of advertising a model whose every request - // the shared KV gate rejects — the v0.7.2 black-hole shape. + // Post-BRIDGE measured-headroom re-guard (v0.7.3, kept): engine + // construction can retain additional load-time memory beyond the + // weights the check above measured. VLM slots reuse the wrapper's + // directly owned text tower; no extraction copy is built. Re-measure + // so a box whose full load-time footprint leaves no serveable KV + // unloads and 503s instead of advertising a model whose every + // request the shared KV gate rejects — the v0.7.2 black-hole shape. // BACKEND-AWARE: a PAGED slot's slabs are committed lazily at // the pool's FIRST ADMISSION (`.atFirstAdmission`, the D1 fix), // NOT at construction — so the measured headroom taken here diff --git a/provider-swift/Sources/ProviderCore/ProviderLoop.swift b/provider-swift/Sources/ProviderCore/ProviderLoop.swift index 4c5e20e4a..875c32fbf 100644 --- a/provider-swift/Sources/ProviderCore/ProviderLoop.swift +++ b/provider-swift/Sources/ProviderCore/ProviderLoop.swift @@ -586,8 +586,8 @@ public actor ProviderLoop { /// the slot never exists (`ensureModelLoaded` unloads + 503s). let engineBundle: ProviderEngineBundle var engineV2: EngineV2Bridge { engineBundle.bridge } - /// Retained for VLM vision preprocessing (the tower shares weights - /// with the extracted text model) and for liveness rebuilds. + /// Retained for VLM vision preprocessing and liveness rebuilds; the + /// wrapper owns the exact text tower retained by the engine. let container: MLXLMCommon.ModelContainer let tokenizer: TokenizerHandle /// Scheduler-free sizing facts (weights, fp16 KV rate, context) — diff --git a/provider-swift/Sources/ProviderCore/Server/StandaloneServer.swift b/provider-swift/Sources/ProviderCore/Server/StandaloneServer.swift index e02cb8a6b..57e0689e9 100644 --- a/provider-swift/Sources/ProviderCore/Server/StandaloneServer.swift +++ b/provider-swift/Sources/ProviderCore/Server/StandaloneServer.swift @@ -118,10 +118,9 @@ private let standaloneLogger = Logger( public actor StandaloneServer { - /// One resident model: its v2 bridge (the serving engine), the loaded - /// container (retained for the VLM media path — the vision tower shares - /// weights with the extracted text model), and the sizing facts the KV - /// re-slice needs. Mirrors `ProviderLoop.ModelSlot`. + /// One resident model: its v2 bridge, loaded container (the VLM owns both + /// vision and the exact text tower served by the bridge), and KV sizing + /// facts. Mirrors `ProviderLoop.ModelSlot`. struct CachedSlot { let bundle: ProviderEngineBundle var bridge: EngineV2Bridge { bundle.bridge } @@ -712,7 +711,6 @@ public actor StandaloneServer { prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: modelId, isVLM: isVLM, - modelDirectory: modelDirectory, container: newcomerBox.borrow(), specDecPreparation: specDecPreparation, assistantLoader: v2TestHooks?.assistantLoader @@ -1371,11 +1369,9 @@ public actor StandaloneServer { var bridge = bundle.bridge // Post-BRIDGE measured-headroom re-guard (mirrors ProviderLoop): - // the engine build (and, for VLM slots, the text-model - // extraction + parity probe) can retain additional load-time - // memory beyond the weights. Re-measure so a box whose full - // load-time footprint leaves no serveable KV tears down instead - // of publishing a model whose every request the KV gate rejects. + // engine construction/JIT may retain load-time memory beyond the + // weights. Re-measure so a box with no serveable KV tears down + // instead of publishing a model whose every request is rejected. // BACKEND-AWARE: a PAGED slot commits only its conservative // physical plan. Require both a useful pool and residual // whole-machine headroom after the build. diff --git a/provider-swift/Sources/ProviderCore/Service/LaunchAgent.swift b/provider-swift/Sources/ProviderCore/Service/LaunchAgent.swift index 4c8bb1f1a..f05ef7170 100644 --- a/provider-swift/Sources/ProviderCore/Service/LaunchAgent.swift +++ b/provider-swift/Sources/ProviderCore/Service/LaunchAgent.swift @@ -107,6 +107,7 @@ public enum LaunchAgent: Sendable { coordinatorURL: String, models: [String] = [], idleTimeout: UInt64? = nil, + configPath: URL? = nil, localEndpoint: LocalEndpointOptions = LocalEndpointOptions() ) throws { // Determine the binary path (current executable) @@ -126,6 +127,7 @@ public enum LaunchAgent: Sendable { coordinatorURL: coordinatorURL, models: models, idleTimeout: idleTimeout, + configPath: configPath, localEndpoint: localEndpoint ) try loadService() @@ -303,6 +305,7 @@ public enum LaunchAgent: Sendable { coordinatorURL: String, models: [String], idleTimeout: UInt64?, + configPath: URL?, localEndpoint: LocalEndpointOptions = LocalEndpointOptions() ) throws { let plist = plistPath() @@ -314,30 +317,14 @@ public enum LaunchAgent: Sendable { let log = logPath().path - // Build the ProgramArguments array. - var programArguments: [String] = [ - binaryPath, - "start", - "--foreground", - "--coordinator-url", - coordinatorURL, - ] - for model in models { - programArguments.append("--model") - programArguments.append(model) - } - if let timeout = idleTimeout { - programArguments.append("--idle-timeout") - programArguments.append("\(timeout)") - } - if localEndpoint.enabled { - programArguments.append("--local-endpoint") - programArguments.append(contentsOf: ["--port", "\(localEndpoint.port)"]) - programArguments.append(contentsOf: ["--bind", localEndpoint.bind]) - if localEndpoint.noAuth { - programArguments.append("--no-auth") - } - } + let programArguments = serviceProgramArguments( + binaryPath: binaryPath, + coordinatorURL: coordinatorURL, + models: models, + idleTimeout: idleTimeout, + configPath: configPath, + localEndpoint: localEndpoint + ) let plistDict = makeServicePlist( label: label, @@ -354,6 +341,44 @@ public enum LaunchAgent: Sendable { try data.write(to: plist, options: .atomic) } + /// Build the child argv without touching launchd or the filesystem. + /// A custom config is explicit so every relaunch reads the same TOML; + /// the canonical default remains implicit and follows normal migration. + static func serviceProgramArguments( + binaryPath: String, + coordinatorURL: String, + models: [String], + idleTimeout: UInt64?, + configPath: URL?, + localEndpoint: LocalEndpointOptions = LocalEndpointOptions() + ) -> [String] { + var arguments = [ + binaryPath, + "start", + "--foreground", + "--coordinator-url", + coordinatorURL, + ] + if let configPath { + arguments.append(contentsOf: ["--config", configPath.standardizedFileURL.path]) + } + for model in models { + arguments.append(contentsOf: ["--model", model]) + } + if let idleTimeout { + arguments.append(contentsOf: ["--idle-timeout", "\(idleTimeout)"]) + } + if localEndpoint.enabled { + arguments.append("--local-endpoint") + arguments.append(contentsOf: ["--port", "\(localEndpoint.port)"]) + arguments.append(contentsOf: ["--bind", localEndpoint.bind]) + if localEndpoint.noAuth { + arguments.append("--no-auth") + } + } + return arguments + } + /// Build the launchd plist dictionary for the provider service. Pure (no I/O) /// so the auto-start and environment-passthrough behavior is unit-testable. /// diff --git a/provider-swift/Sources/ProviderCore/Update/SelfUpdater.swift b/provider-swift/Sources/ProviderCore/Update/SelfUpdater.swift index 78e94b253..0e00b1e80 100644 --- a/provider-swift/Sources/ProviderCore/Update/SelfUpdater.swift +++ b/provider-swift/Sources/ProviderCore/Update/SelfUpdater.swift @@ -1155,11 +1155,19 @@ public struct SelfUpdater: Sendable { + "\(PackagedRuntimeSmoke.mlxLMCommonBundleName)/pagedattention.metal " + "(found \(bundles.count))") } - try BoundedProcess.run( + // The signed child must prove the production TOML projection, + // overwrite precedence, early safe-R1 latch, and packaged AOT before + // it reaches the existing paged-kernel GPU smoke. + let smokeOutput = try BoundedProcess.runCapturingStandardOutput( executable, arguments: ["runtime-smoke"], environment: ["DARKBLOOM_NO_UPDATE_CHECK": "1"], timeout: Self.artifactVerificationTimeout) + guard PackagedRuntimeSmoke.containsGemmaOptimizationSuccessMarker(smokeOutput) + else { + throw UpdateError.replaceFailed( + "packaged runtime smoke omitted the retained Gemma optimization marker") + } } private func verifyCodeSignature( diff --git a/provider-swift/Sources/darkbloom/BenchmarkCommand+Sweep.swift b/provider-swift/Sources/darkbloom/BenchmarkCommand+Sweep.swift index 8eeadeb40..c400483c5 100644 --- a/provider-swift/Sources/darkbloom/BenchmarkCommand+Sweep.swift +++ b/provider-swift/Sources/darkbloom/BenchmarkCommand+Sweep.swift @@ -28,7 +28,8 @@ extension Benchmark { func runThroughputSweep( modelID: String, modelDirectory: URL, - hardware: HardwareInfo + hardware: HardwareInfo, + gemmaOptimizations: GemmaOptimizationSettings ) async throws { let lengths = Self.parsePositiveInts(prefillLengths) guard !lengths.isEmpty else { @@ -66,6 +67,7 @@ extension Benchmark { decodePromptTokens: decodePromptTokens, decodeIterations: decodeIterations, kvBackend: backend, + gemmaOptimizations: gemmaOptimizations, hardware: hardware ) @@ -134,7 +136,8 @@ extension Benchmark { func runSchedulerPrefillBenchmark( modelID: String, - modelDirectory: URL + modelDirectory: URL, + gemmaOptimizations: GemmaOptimizationSettings ) async throws { let lengths = Self.parsePositiveInts(prefillLengths) guard !lengths.isEmpty else { @@ -151,7 +154,8 @@ extension Benchmark { modelDirectory: modelDirectory, promptLengths: lengths, iterations: prefillIterations, - kvBackend: try resolvedKVBackendSelection() + kvBackend: try resolvedKVBackendSelection(), + gemmaOptimizations: gemmaOptimizations ) print(try report.jsonString()) @@ -159,7 +163,8 @@ extension Benchmark { func runArrivalInvarianceBenchmark( modelID: String, - modelDirectory: URL + modelDirectory: URL, + gemmaOptimizations: GemmaOptimizationSettings ) async throws { guard arrivalPromptTokens >= 2 else { printError("--arrival-prompt-tokens must be >= 2") @@ -180,7 +185,8 @@ extension Benchmark { promptTokens: arrivalPromptTokens, decodeTokens: arrivalDecodeTokens, iterations: arrivalIterations, - kvBackend: try resolvedKVBackendSelection() + kvBackend: try resolvedKVBackendSelection(), + gemmaOptimizations: gemmaOptimizations ) print(try report.jsonString()) } diff --git a/provider-swift/Sources/darkbloom/BenchmarkCommand.swift b/provider-swift/Sources/darkbloom/BenchmarkCommand.swift index c2569372a..243233249 100644 --- a/provider-swift/Sources/darkbloom/BenchmarkCommand.swift +++ b/provider-swift/Sources/darkbloom/BenchmarkCommand.swift @@ -1,4 +1,5 @@ import ArgumentParser +import Foundation import ProviderCore import ProviderBenchmark @@ -126,14 +127,47 @@ struct Benchmark: AsyncParsableCommand { var parityPrefixTokens = 28672 mutating func run() async throws { + let snapshot = try loadRuntimeSnapshot( + configPath: configOptions.config, + migrateOnDisk: false) + + // The low-level Gemma controls are process-start latches, so + // `provider.toml` must be projected BEFORE the first MLX device + // access — exactly like the serve path (the shared seam is + // `ServeRuntimePreparer.prepareRuntime`; `Start` forwards to it). + // Otherwise a rollback A/B benchmark silently measures the + // default-enabled stack instead of the configured serving stack. A + // rejected projection aborts before engine construction. + // + // A/B integrity: benchmark artifacts record the process environment + // (scripts/gemma_contbatch/runner.py logs os.environ). A shell-preset + // low-level key that CONFLICTS with the config projection would be + // silently overwritten by apply(), so the artifact metadata would + // then disagree with what was actually measured — refuse to run. + let gemmaSettings = snapshot.config.gemmaOptimizations + if let conflict = ServeRuntimePreparer.conflictingEnvironmentOverride( + settings: gemmaSettings + ) { + printError(""" + benchmark is config-driven: refusing to run with \ + \(conflict.key)=\(conflict.shellValue) set in this shell while \ + provider.toml projects \(conflict.configValue). Toggle with \ + `darkbloom beta` or edit [gemma_optimizations] in provider.toml, \ + then re-run benchmark without the low-level override. + """) + throw ExitCode.failure + } do { - _ = try GPUEnforcement.requireMetal() + try ServeRuntimePreparer.prepareRuntime(settings: gemmaSettings) } catch { printError("\(error)") throw ExitCode.failure } - - let snapshot = try loadRuntimeSnapshot(configOptions: configOptions) + // stderr, not stdout — benchmark subcommands emit machine-parsed JSON + // on stdout (any stray line breaks `darkbloom benchmark`'s consumers). + FileHandle.standardError.write(Data( + "gemma optimizations: prefill_layer18=\(gemmaSettings.prefillLayer18 ? "on" : "off") weighted_r1=\(gemmaSettings.weightedR1 ? "on" : "off")\n" + .utf8)) guard let hardware = snapshot.hardware else { printError("hardware detection failed: \(snapshot.hardwareError?.localizedDescription ?? "unknown")") @@ -167,7 +201,8 @@ struct Benchmark: AsyncParsableCommand { try await runThroughputSweep( modelID: selectedModel.id, modelDirectory: modelPath, - hardware: hardware + hardware: hardware, + gemmaOptimizations: gemmaSettings ) return } @@ -175,7 +210,8 @@ struct Benchmark: AsyncParsableCommand { if schedulerPrefill { try await runSchedulerPrefillBenchmark( modelID: selectedModel.id, - modelDirectory: modelPath + modelDirectory: modelPath, + gemmaOptimizations: gemmaSettings ) return } @@ -183,7 +219,8 @@ struct Benchmark: AsyncParsableCommand { if arrivalInvariance { try await runArrivalInvarianceBenchmark( modelID: selectedModel.id, - modelDirectory: modelPath + modelDirectory: modelPath, + gemmaOptimizations: gemmaSettings ) return } diff --git a/provider-swift/Sources/darkbloom/BetaCommand.swift b/provider-swift/Sources/darkbloom/BetaCommand.swift index 3d32b16a4..330867790 100644 --- a/provider-swift/Sources/darkbloom/BetaCommand.swift +++ b/provider-swift/Sources/darkbloom/BetaCommand.swift @@ -1,16 +1,18 @@ import Foundation import ArgumentParser import ProviderCore +#if canImport(Darwin) +import Darwin +#endif struct Beta: AsyncParsableCommand { static let configuration = CommandConfiguration( commandName: "beta", - abstract: "Manage opt-in beta features.", + abstract: "Manage configurable beta features.", discussion: """ - Beta features are experimental and off by default. Toggling one writes a - field in your provider TOML config, so the change also applies to the - launchd daemon (unlike environment variables, which the daemon does not - inherit). + Beta features are experimental and have feature-specific defaults. + Toggling one writes a field in your provider TOML config, which is the + authority for daemon, `--foreground`, and `--local` processes. Subcommands: list Show all beta features and whether each is on (default). @@ -18,8 +20,9 @@ struct Beta: AsyncParsableCommand { disable Turn a beta feature off. status [feature] Show details for all features, or one. - Most changes require a restart to take effect: - darkbloom beta enable mtp + Changes that affect process-wide optimization state require a restart. + To roll back a default-on feature: + darkbloom beta disable gemma-weighted-r1 darkbloom restart """, subcommands: [List.self, Enable.self, Disable.self, Status.self], @@ -79,7 +82,7 @@ extension Beta { } } print("") - print("Enable with: darkbloom beta enable (then: darkbloom restart)") + print("Change with: darkbloom beta enable|disable (then: darkbloom restart)") print("Details with: darkbloom beta status ") } } @@ -139,7 +142,7 @@ extension Beta { var feature: String mutating func run() async throws { - try setBetaFeature(feature, enabled: true, configOptions: configOptions) + try setBetaFeature(feature, enabled: true, configPath: configOptions.config) } } @@ -154,7 +157,7 @@ extension Beta { var feature: String mutating func run() async throws { - try setBetaFeature(feature, enabled: false, configOptions: configOptions) + try setBetaFeature(feature, enabled: false, configPath: configOptions.config) } } } @@ -168,17 +171,19 @@ private func unknownFeatureError(_ id: String) -> ValidationError { } /// Read-modify-write a single beta feature's config field and persist it. -private func setBetaFeature( +/// +/// Internal (not `private`) so `DarkbloomCLITests` can drive it with temp +/// config fixtures via `configPath`. +func setBetaFeature( _ id: String, enabled: Bool, - configOptions: ConfigOptions + configPath: String? ) throws { guard let feature = BetaFeatures.feature(id: id) else { throw unknownFeatureError(id) } - let snapshot = try loadRuntimeSnapshot(configOptions: configOptions) - var config = snapshot.config + let snapshot = try loadRuntimeSnapshot(configPath: configPath) // Persist to the path the daemon will actually read. With no explicit // --config, loadRuntimeSnapshot may have just migrated a legacy config to @@ -187,26 +192,102 @@ private func setBetaFeature( // the (legacy) snapshot.configPath would leave the restarted daemon on the // stale value. Re-resolving the default returns the post-migration canonical. let savePath: URL - if configOptions.config != nil { + if configPath != nil { savePath = snapshot.configPath } else { savePath = try ConfigManager.defaultConfigPath() } - // Already in the desired state and the target file exists — no-op. - if feature.isEnabled(in: config) == enabled - && FileManager.default.fileExists(atPath: savePath.path) { - print("\(feature.title) (\(feature.id)) is already \(enabled ? "enabled" : "disabled").") - return + // Serialize the load → modify → save window with an exclusive flock, and + // RELOAD inside the lock: the snapshot above and any concurrent `beta` + // process's write can interleave, and using the pre-lock snapshot would be + // a classic lost-update RMW race. + try withExclusiveConfigLock(at: savePath) { + var config: ProviderConfig + if FileManager.default.fileExists(atPath: savePath.path) { + config = try ConfigManager.load(from: savePath) + } else { + config = snapshot.config + } + + // No-op only when the file already PINS the requested value. An absent + // key (or absent [section]) can decode to the same effective value via + // the default, but an explicit enable/disable means "make it so, + // durably" — materialize the key so a future default flip cannot + // silently move this provider. + if feature.isEnabled(in: config) == enabled, + let address = feature.configAddress, + let content = try? String(contentsOf: savePath, encoding: .utf8), + tomlKeyPresent(content, section: address.section, key: address.key) { + print("\(feature.title) (\(feature.id)) is already \(enabled ? "enabled" : "disabled").") + return + } + + feature.apply(enabled, to: &config) + try ConfigManager.save(config, to: savePath) + + print("\(enabled ? "Enabled" : "Disabled") beta feature: \(feature.title) (\(feature.id))") + print(" \(feature.details)") + if feature.requiresRestart { + print(" Restart to apply: darkbloom restart") + } + print(" Config: \(savePath.path)") } +} - feature.apply(enabled, to: &config) - try ConfigManager.save(config, to: savePath) +/// Guards one config-file mutation window with an exclusive `flock(2)` on a +/// stable `.lock` sidecar next to the config file. The lock must +/// NOT be taken out on provider.toml itself: `ConfigManager.save` writes +/// atomically via temp-file + rename, so the config file's inode changes on +/// every save and concurrent writers would be locking DIFFERENT inodes (no +/// mutual exclusion). The sidecar path is never renamed, so every contending +/// process locks the same inode. Closing the fd (the defer) also releases the +/// kernel lock if the explicit LOCK_UN is ever skipped by a throw. +@discardableResult +func withExclusiveConfigLock(at configPath: URL, _ body: () throws -> T) throws -> T { + let directory = configPath.deletingLastPathComponent() + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + let lockURL = directory.appendingPathComponent(configPath.lastPathComponent + ".lock") + + let fd = open(lockURL.path, O_RDWR | O_CREAT, 0o644) + guard fd >= 0 else { + throw ConfigError.writeFailed( + path: lockURL.path, + underlying: NSError(domain: NSPOSIXErrorDomain, code: Int(errno)) + ) + } + defer { close(fd) } + + guard flock(fd, LOCK_EX) == 0 else { + throw ConfigError.writeFailed( + path: lockURL.path, + underlying: NSError(domain: NSPOSIXErrorDomain, code: Int(errno)) + ) + } + defer { _ = flock(fd, LOCK_UN) } - print("\(enabled ? "Enabled" : "Disabled") beta feature: \(feature.title) (\(feature.id))") - print(" \(feature.details)") - if feature.requiresRestart { - print(" Restart to apply: darkbloom restart") + return try body() +} + +/// Whether TOML `content` materially sets `key` inside `[section]`. +/// +/// Line-oriented: tracks the current table header and matches `key = ...` +/// assignments. Only needs to be correct for the flat +/// `[section]\nkey = value` shape `ConfigManager.save` serializes (and that +/// operators hand-edit). A miss here is fail-safe for the caller: unsure +/// means WRITE the key, which is idempotent. +func tomlKeyPresent(_ content: String, section: String, key: String) -> Bool { + var inSection = false + for rawLine in content.split(separator: "\n", omittingEmptySubsequences: false) { + let line = rawLine.trimmingCharacters(in: .whitespaces) + if line.hasPrefix("[") { + inSection = line == "[\(section)]" + continue + } + guard inSection, !line.hasPrefix("#"), + let eqIndex = line.firstIndex(of: "=") else { continue } + let name = line[.. RuntimeSnapshot return try loadRuntimeSnapshot(configPath: configOptions.config) } -func loadRuntimeSnapshot(configPath rawPath: String?) throws -> RuntimeSnapshot { +func loadRuntimeSnapshot( + configPath rawPath: String?, + migrateOnDisk: Bool = true +) throws -> RuntimeSnapshot { let configPath = try resolveConfigPath(rawPath) let configFileExists = FileManager.default.fileExists(atPath: configPath.path) @@ -133,8 +136,11 @@ func loadRuntimeSnapshot(configPath rawPath: String?) throws -> RuntimeSnapshot config = ConfigManager.loadDefault() } - // Auto-migrate stale config values (idempotent, best-effort). - config = migrateConfigIfNeeded(configPath: configPath, config: config) + // Serving/operator commands migrate stale config values. Benchmarking is + // read-only: measurement must never rewrite the input half of an A/B pair. + if migrateOnDisk { + config = migrateConfigIfNeeded(configPath: configPath, config: config) + } let models = hardware.map { ModelScanner.scanModels(hardwareInfo: $0) } ?? [] diff --git a/provider-swift/Sources/darkbloom/RuntimeSmokeCommand.swift b/provider-swift/Sources/darkbloom/RuntimeSmokeCommand.swift index 9086d637b..769529b91 100644 --- a/provider-swift/Sources/darkbloom/RuntimeSmokeCommand.swift +++ b/provider-swift/Sources/darkbloom/RuntimeSmokeCommand.swift @@ -13,6 +13,8 @@ struct RuntimeSmoke: ParsableCommand { var shapes: [String] = [] mutating func run() throws { + try PackagedRuntimeSmoke.verifyGemmaOptimizations() + print(PackagedRuntimeSmoke.gemmaOptimizationSuccessMarker) try PackagedRuntimeSmoke.runPagedKernel(arguments: shapes) print("paged-kernel-runtime-smoke: ok") } diff --git a/provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift b/provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift new file mode 100644 index 000000000..8a9bd4ff5 --- /dev/null +++ b/provider-swift/Sources/darkbloom/ServeRuntimePreparer.swift @@ -0,0 +1,80 @@ +import Foundation +import ProviderCore +#if canImport(Darwin) +import Darwin +#endif + +/// Process-start seam shared by every serving/benchmarking entry point +/// (`start` daemon/foreground/local, `benchmark`): the authoritative +/// `[gemma_optimizations]` TOML projection must be applied to the low-level +/// process environment BEFORE the first MLX device access (`requireMetal`), +/// because those controls are process-start latches in MLX/MLX-LM. +/// +/// Ordering contract: config projection strictly precedes the first MLX +/// touch. A rejected projection throws before `requireMetal()`, so a +/// half-applied weighted-unsort/safe-R1 pair can never reach engine +/// construction. +enum ServeRuntimePreparer { + + /// Apply `settings` to the process environment, then probe Metal. + /// + /// The default closures are the production path; tests replace only the + /// apply/Metal probes so they can assert ordering without constructing an + /// MLX device. + internal static func prepareRuntime( + settings: GemmaOptimizationSettings, + apply: (GemmaOptimizationSettings) throws -> Void = { + try GemmaOptimizationEnvironment.apply($0) + }, + requireMetal: () throws -> Void = { + _ = try GPUEnforcement.requireMetal() + } + ) throws { + try apply(settings) + try requireMetal() + } + + /// One pre-set low-level environment key that CONFLICTS with the config + /// projection a command is about to apply. + struct EnvironmentConflict: Equatable { + /// The low-level environment key (e.g. `MLX_GATHER_QMM_EXPERT_SLICES`). + let key: String + /// The value the operator's shell exported. + let shellValue: String + /// The value `provider.toml` projects (and would overwrite with). + let configValue: String + } + + /// Returns the first pre-existing low-level key whose value DISAGREES with + /// the config projection; nil when every key is unset or already matches. + /// + /// `apply(_:)` overwrites unconditionally (config is authoritative), which + /// is correct for serving — but a benchmark run whose artifact metadata + /// records `os.environ` (scripts/gemma_contbatch/runner.py) would then + /// disagree with what was actually measured. Benchmark-style callers check + /// this first and refuse to run on a conflict instead of silently + /// rewriting the operator's shell. Sorted scan keeps the reported key + /// stable across runs. + internal static func conflictingEnvironmentOverride( + settings: GemmaOptimizationSettings, + getenv: (String) -> String? = { + $0.withCString { Darwin.getenv($0) }.map { String(cString: $0) } + } + ) -> EnvironmentConflict? { + let projection = GemmaOptimizationEnvironment.projection(for: settings) + for key in [ + GemmaOptimizationEnvironment.prefillLayer18Key, + GemmaOptimizationEnvironment.weightedUnsortKey, + GemmaOptimizationEnvironment.safeR1Key, + ].sorted() { + guard let shellValue = getenv(key), + shellValue != projection[key] else { continue } + return EnvironmentConflict( + key: key, + shellValue: shellValue, + configValue: projection[key] ?? "" + ) + } + return nil + } +} diff --git a/provider-swift/Sources/darkbloom/StartCommand+Daemon.swift b/provider-swift/Sources/darkbloom/StartCommand+Daemon.swift index 98e0ed3a5..334cbd6bb 100644 --- a/provider-swift/Sources/darkbloom/StartCommand+Daemon.swift +++ b/provider-swift/Sources/darkbloom/StartCommand+Daemon.swift @@ -13,7 +13,8 @@ extension Start { internal mutating func launchDaemon( snapshot: RuntimeSnapshot, config: ProviderConfig, - coordinatorURL: String + coordinatorURL: String, + configPath: URL? ) async throws { // Run critical checks before downloading models or prompting. try runPreflightChecks(snapshot: snapshot) @@ -44,6 +45,7 @@ extension Start { coordinatorURL: coordinatorURL, models: selectedModelIDs, idleTimeout: idleTimeout ?? (config.backend.idleTimeoutMins > 0 ? config.backend.idleTimeoutMins : nil), + configPath: configPath, localEndpoint: LaunchAgent.LocalEndpointOptions( enabled: localEndpoint, port: port, bind: bind, noAuth: noAuth ) diff --git a/provider-swift/Sources/darkbloom/StartCommand.swift b/provider-swift/Sources/darkbloom/StartCommand.swift index cd500760b..9b934963f 100644 --- a/provider-swift/Sources/darkbloom/StartCommand.swift +++ b/provider-swift/Sources/darkbloom/StartCommand.swift @@ -62,7 +62,6 @@ struct Start: AsyncParsableCommand { mutating func run() async throws { Darkbloom.ensureLogging() - if !foreground { printTermsNotice() } @@ -75,15 +74,6 @@ struct Start: AsyncParsableCommand { throw ExitCode.failure } - // GPU is required. Reject CPU fallback up-front so we never - // come up reporting healthy and then silently churn at 0.5 tok/s. - do { - _ = try GPUEnforcement.requireMetal() - } catch { - printError("\(error)") - throw ExitCode.failure - } - let snapshot = try loadRuntimeSnapshot(configOptions: configOptions) let effectiveCoordinator = coordinatorURL ?? snapshot.config.coordinator.url var effectiveConfig = snapshot.config @@ -91,6 +81,15 @@ struct Start: AsyncParsableCommand { effectiveConfig.backend.idleTimeoutMins = idleTimeout } + // These controls are process-start latches in MLX/MLXLM. Project the + // authoritative TOML before requireMetal() performs the first MLX touch. + do { + try Self.prepareServeRuntime(settings: snapshot.config.gemmaOptimizations) + } catch { + printError("Cannot start: \(error)") + throw ExitCode.failure + } + guard let hardware = snapshot.hardware else { printError("Cannot start: hardware detection failed (\(snapshot.hardwareError?.localizedDescription ?? "unknown"))") throw ExitCode.failure @@ -120,9 +119,32 @@ struct Start: AsyncParsableCommand { try await launchDaemon( snapshot: snapshot, config: effectiveConfig, - coordinatorURL: effectiveCoordinator + coordinatorURL: effectiveCoordinator, + configPath: configOptions.config == nil ? nil : snapshot.configPath ) } } + /// Backward-compatible forwarding shim for the process-start environment + /// projection. The real seam (and its ordering contract: config projection + /// strictly BEFORE the first MLX touch; a rejected projection throws + /// before `requireMetal()`) lives in `ServeRuntimePreparer.prepareRuntime` + /// so `benchmark` mirrors the serve path without referencing `Start`. + /// Tests target `ServeRuntimePreparer` directly. + internal static func prepareServeRuntime( + settings: GemmaOptimizationSettings, + apply: (GemmaOptimizationSettings) throws -> Void = { + try GemmaOptimizationEnvironment.apply($0) + }, + requireMetal: () throws -> Void = { + _ = try GPUEnforcement.requireMetal() + } + ) throws { + try ServeRuntimePreparer.prepareRuntime( + settings: settings, + apply: apply, + requireMetal: requireMetal + ) + } + } diff --git a/provider-swift/Sources/darkbloom/UpdateCommand.swift b/provider-swift/Sources/darkbloom/UpdateCommand.swift index 9bb3b03a7..4c568c2aa 100644 --- a/provider-swift/Sources/darkbloom/UpdateCommand.swift +++ b/provider-swift/Sources/darkbloom/UpdateCommand.swift @@ -1,6 +1,25 @@ import ArgumentParser +import Foundation import ProviderCore +private func updateConfigReadFailureIsMissing(_ error: Error) -> Bool { + let nsError = error as NSError + return (nsError.domain == NSCocoaErrorDomain + && nsError.code == CocoaError.Code.fileNoSuchFile.rawValue) + || (nsError.domain == NSPOSIXErrorDomain + && nsError.code == Int(POSIXErrorCode.ENOENT.rawValue)) +} + +func loadUpdateConfig(configPath: String?) throws -> ProviderConfig { + do { + return try loadRuntimeSnapshot(configPath: configPath).config + } catch ConfigError.readFailed(_, let underlying) + where updateConfigReadFailureIsMissing(underlying) + { + return ConfigManager.loadDefault() + } +} + struct Update: AsyncParsableCommand { static let configuration = CommandConfiguration( commandName: "update", @@ -22,13 +41,7 @@ struct Update: AsyncParsableCommand { var overrideQuarantine = false mutating func run() async throws { - let config: ProviderConfig - do { - let snapshot = try loadRuntimeSnapshot(configOptions: configOptions) - config = snapshot.config - } catch { - config = ConfigManager.loadDefault() - } + let config = try loadUpdateConfig(configPath: configOptions.config) print("darkbloom update") print("Current version: \(ProviderCore.version)") diff --git a/provider-swift/Tests/DarkbloomCLITests/BenchmarkSweepExitTests.swift b/provider-swift/Tests/DarkbloomCLITests/BenchmarkSweepExitTests.swift index c8c618f5f..4af107b31 100644 --- a/provider-swift/Tests/DarkbloomCLITests/BenchmarkSweepExitTests.swift +++ b/provider-swift/Tests/DarkbloomCLITests/BenchmarkSweepExitTests.swift @@ -129,10 +129,10 @@ struct BenchmarkSweepExitTests { // Under `auto` the operator named no backend, so a cell that could // not build one is an ordinary bad run, not a broken promise — the // same asymmetry `EngineV2KVBackendPolicy.degradesPagedFailure` - // encodes for the engine. `auto` resolves paged but DEGRADES on a - // paged failure rather than refusing, so a paged capacity refusal is - // still not reachable here; what is reachable is an ordinary - // construction error, and it keeps its exit status. + // encodes for the engine. `auto` resolves contiguous as of v0.8.1, + // so a paged capacity refusal is not reachable here; what is reachable + // is an ordinary construction error, and it keeps its exit status. The + // degrade rule remains for any future release that resolves auto paged. // Nil message ⇒ `runThroughputSweep` returns normally ⇒ 0. #expect(Benchmark.sweepFailureMessage( backend: .auto, failure: nil, coverage: coverage()) == nil) diff --git a/provider-swift/Tests/DarkbloomCLITests/BetaCommandTests.swift b/provider-swift/Tests/DarkbloomCLITests/BetaCommandTests.swift new file mode 100644 index 000000000..4d66f2a9b --- /dev/null +++ b/provider-swift/Tests/DarkbloomCLITests/BetaCommandTests.swift @@ -0,0 +1,239 @@ +import Darwin +import Foundation +import ProviderCore +import Testing + +@testable import darkbloom + +@Suite("Beta command config mutation") +struct BetaCommandTests { + + /// Write `toml` (when non-nil) into a unique temp `provider.toml` and + /// return its URL. `~/.config/darkbloom/provider.toml` is guarded around + /// every toggle call: a non-canonical config path triggers + /// loadRuntimeSnapshot's legacy→canonical copy when the canonical file is + /// ABSENT, which would otherwise plant test fixtures in the operator's + /// real config on a fresh machine. + private func makeTempConfig(_ toml: String?) throws -> URL { + let directory = FileManager.default.temporaryDirectory + .appendingPathComponent("beta-cfg-\(UUID().uuidString)") + try FileManager.default.createDirectory( + at: directory, withIntermediateDirectories: true) + let url = directory.appendingPathComponent("provider.toml") + if let toml { + try toml.write(to: url, atomically: true, encoding: .utf8) + } + return url + } + + private func withGuardedCanonicalConfig( + _ body: () throws -> Void + ) rethrows { + let canonical = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".config/darkbloom/provider.toml") + let existedBefore = FileManager.default.fileExists(atPath: canonical.path) + defer { + if !existedBefore, + FileManager.default.fileExists(atPath: canonical.path) { + try? FileManager.default.removeItem(at: canonical) + } + } + try body() + } + + @Test("enable with an absent section writes the key despite the matching default") + func enableMaterializesAbsentKey() throws { + let url = try makeTempConfig(""" + config_version = 2 + + [provider] + name = "beta-test" + """) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + + // weightedR1 already decodes to true via the default; the old code + // no-oped ("already enabled") without pinning anything. + try withGuardedCanonicalConfig { + try setBetaFeature("gemma-weighted-r1", enabled: true, configPath: url.path) + } + + let written = try String(contentsOf: url, encoding: .utf8) + #expect(written.contains("[gemma_optimizations]")) + #expect(written.contains("weighted_r1 = true")) + let reloaded = try ConfigManager.load(from: url) + #expect(reloaded.gemmaOptimizations.weightedR1) + } + + @Test("enable after the materializing write is a true no-op") + func secondEnableDoesNotRewrite() throws { + let url = try makeTempConfig(""" + [provider] + name = "beta-test" + """) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + + try withGuardedCanonicalConfig { + try setBetaFeature("gemma-weighted-r1", enabled: true, configPath: url.path) + let pinned = try String(contentsOf: url, encoding: .utf8) + + try setBetaFeature("gemma-weighted-r1", enabled: true, configPath: url.path) + let after = try String(contentsOf: url, encoding: .utf8) + + #expect(after == pinned) + } + } + + @Test("a key pinned at the target value is a no-op without a rewrite") + func pinnedKeyIsNoOp() throws { + let url = try makeTempConfig(""" + config_version = 2 + + [provider] + name = "beta-test" + + [gemma_optimizations] + weighted_r1 = true + """) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + let before = try String(contentsOf: url, encoding: .utf8) + + try withGuardedCanonicalConfig { + try setBetaFeature("gemma-weighted-r1", enabled: true, configPath: url.path) + } + + let after = try String(contentsOf: url, encoding: .utf8) + #expect(after == before) + } + + @Test("disable with an absent key still writes a default-off feature") + func disableMaterializesAbsentKey() throws { + // MTP defaults off: disabling an absent key used to no-op without + // persisting the operator's intent. + let url = try makeTempConfig(""" + [provider] + name = "beta-test" + """) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + + try withGuardedCanonicalConfig { + try setBetaFeature("mtp", enabled: false, configPath: url.path) + } + + let written = try String(contentsOf: url, encoding: .utf8) + #expect(written.contains("mtp = false")) + #expect(tomlKeyPresent(written, section: "backend", key: "mtp")) + } + + @Test("disable flips a pinned key and keeps its neighbour") + func disableFlipsPinnedKey() throws { + let url = try makeTempConfig(""" + [provider] + name = "beta-test" + + [gemma_optimizations] + prefill_layer18 = false + """) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + + try withGuardedCanonicalConfig { + try setBetaFeature("gemma-weighted-r1", enabled: false, configPath: url.path) + } + + let reloaded = try ConfigManager.load(from: url) + #expect(!reloaded.gemmaOptimizations.weightedR1) + #expect(!reloaded.gemmaOptimizations.prefillLayer18) + let written = try String(contentsOf: url, encoding: .utf8) + #expect(written.contains("weighted_r1 = false")) + #expect(written.contains("prefill_layer18 = false")) + } + + @Test("an unknown feature id is rejected before touching the file") + func unknownFeatureThrows() throws { + let url = try makeTempConfig(""" + [provider] + name = "beta-test" + """) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + let before = try String(contentsOf: url, encoding: .utf8) + + do { + try setBetaFeature("gemma-expert-packing", enabled: true, configPath: url.path) + Issue.record("gemma-expert-packing is not a beta feature in this build") + } catch { + // ValidationError naming the known feature ids. + } + + let after = try String(contentsOf: url, encoding: .utf8) + #expect(after == before) + } + + @Test("enabling into a missing config file creates it with the key") + func missingFileIsWritten() throws { + let url = try makeTempConfig(nil) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + + try withGuardedCanonicalConfig { + try setBetaFeature("gemma-prefill-layer18", enabled: false, configPath: url.path) + } + + let written = try String(contentsOf: url, encoding: .utf8) + #expect(written.contains("prefill_layer18 = false")) + } + + // MARK: - tomlKeyPresent + + @Test("tomlKeyPresent matches keys inside their own section only") + func tomlKeyPresentSectioning() { + let content = """ + [provider] + name = "x" + + [gemma_optimizations] + weighted_r1 = false + + [backend] + kv_quant = true + """ + + #expect(tomlKeyPresent(content, section: "gemma_optimizations", key: "weighted_r1")) + #expect(tomlKeyPresent(content, section: "backend", key: "kv_quant")) + #expect(!tomlKeyPresent(content, section: "backend", key: "weighted_r1")) + #expect(!tomlKeyPresent(content, section: "gemma_optimizations", key: "kv_quant")) + #expect(!tomlKeyPresent(content, section: "gemma_optimizations", key: "prefill_layer18")) + } + + @Test("tomlKeyPresent ignores commented-out keys") + func tomlKeyPresentIgnoresComments() { + let content = """ + [gemma_optimizations] + # weighted_r1 = false + """ + #expect(!tomlKeyPresent(content, section: "gemma_optimizations", key: "weighted_r1")) + } + + // MARK: - config flock + + @Test("the config lock excludes a second exclusive flock on the sidecar") + func configLockMutualExclusion() throws { + let url = try makeTempConfig(nil) + defer { try? FileManager.default.removeItem(at: url.deletingLastPathComponent()) } + + let lockPath = url.path + ".lock" + try withExclusiveConfigLock(at: url) { + let fd = open(lockPath, O_RDWR | O_CREAT, 0o644) + #expect(fd >= 0) + defer { close(fd) } + // flock is per open-file-description: a second descriptor to the + // same sidecar must fail LOCK_EX|LOCK_NB while the guard holds it. + #expect(flock(fd, LOCK_EX | LOCK_NB) != 0) + #expect(errno == EWOULDBLOCK) + } + + // After the guard released it, an exclusive lock succeeds again. + let fd = open(lockPath, O_RDWR | O_CREAT, 0o644) + #expect(fd >= 0) + defer { close(fd) } + #expect(flock(fd, LOCK_EX | LOCK_NB) == 0) + _ = flock(fd, LOCK_UN) + } +} diff --git a/provider-swift/Tests/DarkbloomCLITests/RuntimeSnapshotConfigTests.swift b/provider-swift/Tests/DarkbloomCLITests/RuntimeSnapshotConfigTests.swift new file mode 100644 index 000000000..448a0234d --- /dev/null +++ b/provider-swift/Tests/DarkbloomCLITests/RuntimeSnapshotConfigTests.swift @@ -0,0 +1,96 @@ +import Foundation +import ProviderCore +import Testing + +@testable import darkbloom + +/// Operator-facing pin for the v0.8.2 config-loading contract: every +/// snapshot-based command (`start`, `benchmark`, `beta`, ...) loads its config +/// through `loadRuntimeSnapshot`. A MISSING file must keep defaulting +/// leniently; an EXISTING but malformed file must fail loudly instead of +/// silently re-enabling whole-config defaults (e.g. re-arming the default-on +/// Gemma stack after a botched rollback edit). +@Suite("Runtime snapshot config loading") +struct RuntimeSnapshotConfigTests { + + private func tempConfigURL() -> URL { + FileManager.default.temporaryDirectory + .appendingPathComponent("snapshot-cfg-\(UUID().uuidString).toml") + } + + @Test("a missing config file keeps the lenient default path") + func missingConfigYieldsDefaults() throws { + let missing = tempConfigURL() + + let snapshot = try loadRuntimeSnapshot(configPath: missing.path) + + #expect(!snapshot.configFileExists) + // Missing section/keys decode defaulted: the optimisation stack stays + // on; hardware-derived name when detection succeeds, "darkbloom" else. + #expect(snapshot.config.gemmaOptimizations == GemmaOptimizationSettings()) + #expect(snapshot.config.gemmaOptimizations.prefillLayer18) + #expect(snapshot.config.gemmaOptimizations.weightedR1) + } + + @Test("a malformed config file fails the load loudly") + func malformedConfigThrows() throws { + let url = tempConfigURL() + try """ + [provider] + name = "snapshot-test" + + [gemma_optimizations] + weighted_r1 = 0 + """.write(to: url, atomically: true, encoding: .utf8) + defer { try? FileManager.default.removeItem(at: url) } + + do { + _ = try loadRuntimeSnapshot(configPath: url.path) + Issue.record("a malformed provider.toml must abort, not silently default") + } catch ConfigError.parseFailed(let detail) { + #expect(!detail.isEmpty) + } + } + + @Test("benchmark-style loading never rewrites its config") + func readOnlyLoadPreservesFile() throws { + let url = tempConfigURL() + let original = """ + config_version = 1 + + [provider] + name = "benchmark-a-b" + + [backend] + engine_v2_max_concurrent = 8 + + [gemma_optimizations] + prefill_layer18 = false + weighted_r1 = false + """ + try original.write(to: url, atomically: true, encoding: .utf8) + defer { try? FileManager.default.removeItem(at: url) } + + let snapshot = try loadRuntimeSnapshot( + configPath: url.path, + migrateOnDisk: false) + + #expect(!snapshot.config.gemmaOptimizations.prefillLayer18) + #expect(!snapshot.config.gemmaOptimizations.weightedR1) + #expect(try String(contentsOf: url, encoding: .utf8) == original) + } + + @Test("update propagates existing-path read failures") + func updateRejectsUnreadableExistingPath() throws { + let url = tempConfigURL() + try FileManager.default.createDirectory(at: url, withIntermediateDirectories: false) + defer { try? FileManager.default.removeItem(at: url) } + + do { + _ = try loadUpdateConfig(configPath: url.path) + Issue.record("an existing path that cannot be read as TOML must not default") + } catch ConfigError.readFailed(let path, _) { + #expect(path == url.path) + } + } +} diff --git a/provider-swift/Tests/DarkbloomCLITests/StartCommandTests.swift b/provider-swift/Tests/DarkbloomCLITests/StartCommandTests.swift new file mode 100644 index 000000000..294bfbe80 --- /dev/null +++ b/provider-swift/Tests/DarkbloomCLITests/StartCommandTests.swift @@ -0,0 +1,177 @@ +import ArgumentParser +import Darwin +import ProviderCore +import Testing + +@testable import darkbloom + +@Suite("Serve runtime preparation (shared Start/Benchmark seam)") +struct StartCommandTests { + @Test("TOML projection precedes the first MLX touch") + func projectionPrecedesMetal() throws { + let settings = GemmaOptimizationSettings( + prefillLayer18: false, + weightedR1: false + ) + var events: [String] = [] + + try ServeRuntimePreparer.prepareRuntime( + settings: settings, + apply: { received in + #expect(received.prefillLayer18 == false) + #expect(received.weightedR1 == false) + events.append("projection") + }, + requireMetal: { + events.append("metal") + } + ) + + #expect(events == ["projection", "metal"]) + } + + @Test("a rejected projection aborts before the first MLX touch") + func rejectedProjectionSkipsMetal() { + var events: [String] = [] + let failure = GemmaOptimizationEnvironment.ApplicationFailure( + keys: [ + GemmaOptimizationEnvironment.safeR1Key, + GemmaOptimizationEnvironment.weightedUnsortKey, + ], + code: ENOMEM + ) + + // The coupled weighted-unsort/safe-R1 pair is a process-start latch: + // a half-applied projection must never reach engine construction. + do { + try ServeRuntimePreparer.prepareRuntime( + settings: GemmaOptimizationSettings(), + apply: { _ in + events.append("projection") + throw failure + }, + requireMetal: { + events.append("metal") + } + ) + Issue.record("a rejected projection must abort serve preparation") + } catch let error as GemmaOptimizationEnvironment.ApplicationFailure { + #expect(error == failure) + #expect("\(error)" == failure.description) + } catch { + Issue.record("expected ApplicationFailure, got \(error)") + } + + #expect(events == ["projection"]) + } + + @Test("the default projection path is the environment apply boundary") + func defaultApplyProjectsSettings() throws { + let settings = GemmaOptimizationSettings( + prefillLayer18: false, + weightedR1: true + ) + let projection = GemmaOptimizationEnvironment.projection(for: settings) + let saved = projection.keys.reduce(into: [String: String?]()) { out, key in + out[key] = key.withCString { getenv($0).map { String(cString: $0) } } + } + defer { + for (key, value) in saved { + if let value { + _ = setenv(key, value, 1) + } else { + _ = unsetenv(key) + } + } + } + var metalProbed = false + + try ServeRuntimePreparer.prepareRuntime( + settings: settings, + requireMetal: { metalProbed = true } + ) + + for (key, value) in projection { + let observed = key.withCString { getenv($0).map { String(cString: $0) } } + #expect(observed == value) + } + #expect(metalProbed) + } + + @Test("the Start compatibility shim forwards to the shared seam") + func startShimForwards() throws { + var events: [String] = [] + try Start.prepareServeRuntime( + settings: GemmaOptimizationSettings(), + apply: { _ in events.append("projection") }, + requireMetal: { events.append("metal") } + ) + #expect(events == ["projection", "metal"]) + } + + @Test("benchmark env guard: a conflicting shell preset is rejected") + func conflictingEnvironmentOverrideReportsConflict() throws { + let settings = GemmaOptimizationSettings( + prefillLayer18: true, + weightedR1: true + ) + let conflict = ServeRuntimePreparer.conflictingEnvironmentOverride( + settings: settings + ) { key in + // Operator rolled back via the shell, config still selects on. + key == GemmaOptimizationEnvironment.safeR1Key ? "0" : nil + } + + let found = try #require(conflict) + #expect(found.key == GemmaOptimizationEnvironment.safeR1Key) + #expect(found.shellValue == "0") + #expect(found.configValue == "1") + } + + @Test("benchmark env guard: the paired weighted-unsort key is checked too") + func conflictingEnvironmentOverrideChecksWeightedKey() { + let settings = GemmaOptimizationSettings( + prefillLayer18: false, + weightedR1: false + ) + let conflict = ServeRuntimePreparer.conflictingEnvironmentOverride( + settings: settings + ) { key in + key == GemmaOptimizationEnvironment.weightedUnsortKey ? "1" : nil + } + + #expect(conflict?.key == GemmaOptimizationEnvironment.weightedUnsortKey) + #expect(conflict?.shellValue == "1") + #expect(conflict?.configValue == "0") + } + + @Test("benchmark env guard: matching or unset shell values are not flagged") + func conflictingEnvironmentOverrideAllowsConsistent() { + let settings = GemmaOptimizationSettings( + prefillLayer18: true, + weightedR1: false + ) + let projection = GemmaOptimizationEnvironment.projection(for: settings) + + // Every preset value EQUALS the projection — consistent, proceed. + let consistent = ServeRuntimePreparer.conflictingEnvironmentOverride( + settings: settings + ) { projection[$0] } + #expect(consistent == nil) + + // Nothing preset at all — proceed. + let unset = ServeRuntimePreparer.conflictingEnvironmentOverride( + settings: settings + ) { _ in nil } + #expect(unset == nil) + } + + @Test("start accepts an explicit custom config") + func customConfigParses() throws { + let parsed = try Darkbloom.parseAsRoot([ + "start", "--config", "/tmp/custom provider.toml", "--foreground", + ]) + let command = try #require(parsed as? Start) + #expect(command.configOptions.config == "/tmp/custom provider.toml") + } +} diff --git a/provider-swift/Tests/ProviderCoreTests/BetaFeaturesTests.swift b/provider-swift/Tests/ProviderCoreTests/BetaFeaturesTests.swift index f9dfee2a4..574ac17d7 100644 --- a/provider-swift/Tests/ProviderCoreTests/BetaFeaturesTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/BetaFeaturesTests.swift @@ -12,9 +12,13 @@ struct BetaFeaturesTests { ) } - @Test("registry exposes beta features") + @Test("registry exposes current features") func registryContainsExpectedFeatures() { + #expect(BetaFeatures.all.contains { $0.id == "gemma-prefill-layer18" }) + #expect(BetaFeatures.all.contains { $0.id == "gemma-weighted-r1" }) #expect(BetaFeatures.all.contains { $0.id == "mtp" }) + #expect(!BetaFeatures.all.contains { $0.id == "gemma-expert-packing" }) + #expect(!BetaFeatures.all.contains { $0.id == "gemma-dense-packing" }) // adaptive-prefill was retired with the legacy engine (v0.7.5); // kv-quant was retired with KV quantization itself (v0.8.0). #expect(!BetaFeatures.all.contains { $0.id == "adaptive-prefill" }) @@ -23,6 +27,11 @@ struct BetaFeaturesTests { @Test("feature lookup is case-insensitive and nil for unknown ids") func featureLookup() { + #expect( + BetaFeatures.feature(id: "GEMMA-PREFILL-LAYER18")?.id + == "gemma-prefill-layer18" + ) + #expect(BetaFeatures.feature(id: "Gemma-Weighted-R1")?.id == "gemma-weighted-r1") #expect(BetaFeatures.feature(id: "mtp")?.id == "mtp") #expect(BetaFeatures.feature(id: "MTP")?.id == "mtp") #expect(BetaFeatures.feature(id: "ADAPTIVE-PREFILL") == nil) // retired v0.7.5 @@ -37,6 +46,35 @@ struct BetaFeaturesTests { #expect(feature.requiresRestart == true) } + @Test("retained Gemma controls default on and require restart") + func gemmaControlsDefaultOn() { + let config = freshConfig() + for id in ["gemma-prefill-layer18", "gemma-weighted-r1"] { + let feature = BetaFeatures.feature(id: id)! + #expect(feature.isEnabled(in: config)) + #expect(feature.requiresRestart) + } + } + + @Test("each retained Gemma row mutates exactly one setting") + func gemmaRowsAreScoped() { + let layer = BetaFeatures.feature(id: "gemma-prefill-layer18")! + let weighted = BetaFeatures.feature(id: "gemma-weighted-r1")! + var config = freshConfig() + + layer.apply(false, to: &config) + #expect(!config.gemmaOptimizations.prefillLayer18) + #expect(config.gemmaOptimizations.weightedR1) + + weighted.apply(false, to: &config) + #expect(!config.gemmaOptimizations.prefillLayer18) + #expect(!config.gemmaOptimizations.weightedR1) + + weighted.apply(true, to: &config) + #expect(!config.gemmaOptimizations.prefillLayer18) + #expect(config.gemmaOptimizations.weightedR1) + } + @Test("apply toggles the backing config field both ways") func applyTogglesField() { let feature = BetaFeatures.feature(id: "mtp")! @@ -69,15 +107,21 @@ struct BetaFeaturesTests { #expect(config.backend.mtpDrafterPath == before.backend.mtpDrafterPath) #expect(config.provider == before.provider) #expect(config.coordinator == before.coordinator) + #expect(config.gemmaOptimizations == before.gemmaOptimizations) } - @Test("enabledIDs reflects the current config") + @Test("enabledIDs reflects default-on and explicit beta settings") func enabledIDsReflectsConfig() { var config = freshConfig() - #expect(BetaFeatures.enabledIDs(in: config).isEmpty) + #expect(BetaFeatures.enabledIDs(in: config) == [ + "gemma-prefill-layer18", "gemma-weighted-r1", + ]) + BetaFeatures.feature(id: "gemma-weighted-r1")!.apply(false, to: &config) BetaFeatures.feature(id: "mtp")!.apply(true, to: &config) - #expect(BetaFeatures.enabledIDs(in: config) == ["mtp"]) + #expect(BetaFeatures.enabledIDs(in: config) == [ + "gemma-prefill-layer18", "mtp", + ]) } @Test("toggling mtp survives a TOML round-trip") @@ -93,4 +137,19 @@ struct BetaFeaturesTests { #expect(feature.isEnabled(in: decoded) == true) } + @Test("retained Gemma beta toggles survive a TOML round trip") + func gemmaRowsRoundTripThroughTOML() { + var config = freshConfig() + BetaFeatures.feature(id: "gemma-prefill-layer18")!.apply(false, to: &config) + BetaFeatures.feature(id: "gemma-weighted-r1")!.apply(false, to: &config) + + let toml = ConfigManager.serialize(config) + let decoded = ConfigManager.parse(toml) + + #expect(toml.contains("prefill_layer18")) + #expect(toml.contains("weighted_r1")) + #expect(!decoded.gemmaOptimizations.prefillLayer18) + #expect(!decoded.gemmaOptimizations.weightedR1) + } + } diff --git a/provider-swift/Tests/ProviderCoreTests/ConfigValidationTests.swift b/provider-swift/Tests/ProviderCoreTests/ConfigValidationTests.swift new file mode 100644 index 000000000..2c4a840f9 --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/ConfigValidationTests.swift @@ -0,0 +1,159 @@ +import Foundation +import Testing +@testable import ProviderCore + +/// v0.8.2 config-loading hardening: the lenient `ConfigManager.parse` stays +/// for tests only. The production FILE path (`load(from:)` → +/// `parseValidating`) must fail loudly — a malformed `[gemma_optimizations]` +/// entry previously decoded to whole-config defaults and silently re-enabled +/// the default-on stack (and silently applied any other default) with no log. +@Suite("Config strict loading") +struct ConfigValidationTests { + + private func writeTempConfig(_ toml: String) throws -> URL { + let url = FileManager.default.temporaryDirectory + .appendingPathComponent("config-validation-\(UUID().uuidString).toml") + try toml.write(to: url, atomically: true, encoding: .utf8) + return url + } + + // MARK: - parseValidating + + @Test("parseValidating decodes a well-formed config") + func parseValidatingDecodesValid() throws { + let config = try ConfigManager.parseValidating(""" + [provider] + name = "strict-provider" + + [gemma_optimizations] + prefill_layer18 = false + weighted_r1 = false + """) + + #expect(config.provider.name == "strict-provider") + #expect(!config.gemmaOptimizations.prefillLayer18) + #expect(!config.gemmaOptimizations.weightedR1) + } + + @Test("parseValidating still defaults a missing optimisation section on") + func parseValidatingDefaultsMissingSectionOn() throws { + let config = try ConfigManager.parseValidating(""" + [provider] + name = "strict-provider" + """) + + #expect(config.gemmaOptimizations == GemmaOptimizationSettings()) + #expect(config.gemmaOptimizations.prefillLayer18) + #expect(config.gemmaOptimizations.weightedR1) + } + + @Test("parseValidating keys off per-key defaults for a partial section") + func parseValidatingDefaultsMissingKeysOn() throws { + let config = try ConfigManager.parseValidating(""" + [gemma_optimizations] + weighted_r1 = false + """) + + #expect(config.gemmaOptimizations.prefillLayer18) + #expect(!config.gemmaOptimizations.weightedR1) + } + + @Test("parseValidating rejects a mis-typed rollback key") + func parseValidatingRejectsIntegerBool() throws { + do { + _ = try ConfigManager.parseValidating(""" + [gemma_optimizations] + weighted_r1 = 0 + """) + Issue.record("weighted_r1 = 0 is not a Bool and must not decode") + } catch ConfigError.parseFailed(let detail) { + #expect(!detail.isEmpty) + } + } + + @Test("parseValidating rejects a quoted rollback key") + func parseValidatingRejectsQuotedBool() throws { + do { + _ = try ConfigManager.parseValidating(""" + [gemma_optimizations] + prefill_layer18 = "false" + """) + Issue.record("prefill_layer18 = \"false\" is not a Bool and must not decode") + } catch ConfigError.parseFailed(let detail) { + #expect(!detail.isEmpty) + } + } + + @Test("parseValidating rejects syntactically broken TOML") + func parseValidatingRejectsBrokenSyntax() throws { + do { + _ = try ConfigManager.parseValidating("[gemma_optimizations\nbroken") + Issue.record("syntactically broken TOML must not decode") + } catch ConfigError.parseFailed(let detail) { + #expect(!detail.isEmpty) + } + } + + // MARK: - file-loading boundary + + @Test("a malformed existing config file fails loading loudly") + func malformedFileThrowsOnLoad() throws { + let url = try writeTempConfig(""" + [provider] + name = "strict-provider" + + [gemma_optimizations] + weighted_r1 = 0 + """) + defer { try? FileManager.default.removeItem(at: url) } + + do { + _ = try ConfigManager.load(from: url) + Issue.record("loading a malformed config file must throw, not default") + } catch ConfigError.parseFailed(let detail) { + #expect(!detail.isEmpty) + } + } + + @Test("a well-formed existing config file still loads") + func wellFormedFileLoads() throws { + let url = try writeTempConfig(""" + [provider] + name = "strict-provider" + + [gemma_optimizations] + prefill_layer18 = false + """) + defer { try? FileManager.default.removeItem(at: url) } + + let config = try ConfigManager.load(from: url) + #expect(config.provider.name == "strict-provider") + #expect(!config.gemmaOptimizations.prefillLayer18) + #expect(config.gemmaOptimizations.weightedR1) + } + + @Test("a missing config file reports a read failure, not a parse failure") + func missingFileIsReadFailure() throws { + let missing = FileManager.default.temporaryDirectory + .appendingPathComponent("config-validation-missing-\(UUID().uuidString).toml") + + do { + _ = try ConfigManager.load(from: missing) + Issue.record("loading a missing file must throw") + } catch ConfigError.readFailed { + // Expected: missing-file LENIENCY lives one layer up + // (loadDefault / loadRuntimeSnapshot substitute defaults there), + // never a silent whole-config default from this boundary. + } + } + + @Test("the lenient test-facing parse keeps its historical contract") + func lenientParseStillDefaultsOnMalformed() { + let config = ConfigManager.parse(""" + [gemma_optimizations] + weighted_r1 = 0 + """) + + #expect(config.gemmaOptimizations == GemmaOptimizationSettings()) + } +} diff --git a/provider-swift/Tests/ProviderCoreTests/EngineV2BridgeTests.swift b/provider-swift/Tests/ProviderCoreTests/EngineV2BridgeTests.swift index 372a8f25d..35e317d2a 100644 --- a/provider-swift/Tests/ProviderCoreTests/EngineV2BridgeTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/EngineV2BridgeTests.swift @@ -1357,8 +1357,6 @@ struct EngineV2FailLoudFactoryTests { EngineV2ProductionError.noKVHeadroom) == .noKVHeadroom) #expect(EngineV2RefusalReason.classify( EngineV2ProductionError.unsupportedModel("Qwen3Model")) == .unsupportedModel) - #expect(EngineV2RefusalReason.classify( - EngineV2VLMTextExtractionError.parityMismatch("x")) == .vlmExtractionFailed) #expect(EngineV2RefusalReason.classify(SomeError()) == .engineInitFailed) } @@ -1442,7 +1440,6 @@ struct EngineV2FailLoudFactoryTests { let cases: [(any Error, String)] = [ (EngineV2ProductionError.noKVHeadroom, "no_kv_headroom"), (EngineV2ProductionError.unsupportedModel("StubModel"), "unsupported_model"), - (EngineV2VLMTextExtractionError.parityMismatch("probe"), "vlm_extraction_failed"), ] for (error, expectedReason) in cases { let telemetry = TelemetrySink() diff --git a/provider-swift/Tests/ProviderCoreTests/EngineV2KVBackendGateTests.swift b/provider-swift/Tests/ProviderCoreTests/EngineV2KVBackendGateTests.swift index f967b34c3..5e37a05ad 100644 --- a/provider-swift/Tests/ProviderCoreTests/EngineV2KVBackendGateTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/EngineV2KVBackendGateTests.swift @@ -138,9 +138,9 @@ private let gateTestCapacity = 8 << 20 // 8 MiB pool — tiny but constructible /// Hermetic default for every construction in this suite: point the /// crash-loop guard store at a file that can never decode. Without it a /// developer box whose REAL provider tripped the guard on the checked-out -/// version would fail every `.auto`-resolves-paged assertion here — the -/// same reason tests inject `environment:` instead of inheriting the -/// shell's kill switch. A caller's explicit value wins. +/// version would fail every explicit-paged assertion here — the same reason +/// tests inject `environment:` instead of inheriting the shell's kill switch. +/// A caller's explicit value wins. private let hermeticGuardEnvironment = [KVBackendGuardStore.pathEnvKey: "/dev/null"] private func gateEnvironment(_ overrides: [String: String] = [:]) -> [String: String] { @@ -158,8 +158,9 @@ private func makeBuild( model: model, tokenizer: StubBridgeTokenizer(), kvBytesCapacity: gateTestCapacity, - // Deliberately 2, not the production 8: these gates assert BACKEND - // SELECTION, and a small pool keeps construction cheap. + // Deliberately 2: these gates assert BACKEND SELECTION, and a small + // pool keeps construction cheap. Production defaults to B=4 while + // still supporting explicit overrides through B=8. maxConcurrentRequests: 2, prefixCache: nil, kvBackend: kvBackend, @@ -345,15 +346,13 @@ struct EngineV2KVBackendGateTests { #expect(reason?.hasPrefix("physical_capacity:") == true) } - @Test("`.auto` still degrades when paged cannot be served") - func autoDegradesOnPagedFailure() async throws { - // Layer 5's degrade half, as a predicate. Since v0.8.0 `.auto` - // RESOLVES PAGED, so the real construction path reaches this branch - // on every paged-ineligible box in the fleet — it is the COMMON - // path, not a hypothetical. The predicate stays pinned here; the - // three construction tests below drive the same answer through the - // real factory, one per failure stage (preflight, capacity - // planning, pool construction). + @Test("non-explicit selections permit paged-failure degradation") + func nonExplicitSelectionsPermitPagedFailureDegradation() async throws { + // Layer 5's degrade half, pinned as a policy predicate. `.auto` + // resolves contiguous as of v0.8.1 and therefore does not currently + // enter a paged-failure branch. The predicate remains the fail-open + // contract if a future release selects paged automatically; an + // explicit `.paged` request must continue to refuse instead. #expect(EngineV2KVBackendPolicy.degradesPagedFailure(selection: .auto)) #expect(EngineV2KVBackendPolicy.degradesPagedFailure(selection: .contiguous)) #expect(!EngineV2KVBackendPolicy.degradesPagedFailure(selection: .paged)) @@ -1005,9 +1004,9 @@ struct EngineV2KVBackendGateTests { /// resolved, so the next capability cannot land with the same invisible /// gap. /// - /// `preparedModel` is supplied so the VLM text extraction (which needs a - /// checkpoint directory) is skipped: the subject here is the ROUTING for - /// `isVLM: true`, not the extraction. + /// `preparedModel` is supplied so real VLM wrapper resolution is skipped: + /// the subject here is the ROUTING for `isVLM: true`, not selection of the + /// wrapper's directly owned text tower. @Test("slot factory routes a VLM slot to paged now that the cache vouches") func slotFactoryRoutesVLMToPagedWhenTheCacheVouches() async throws { #expect( diff --git a/provider-swift/Tests/ProviderCoreTests/EngineV2LivenessRecoveryTests.swift b/provider-swift/Tests/ProviderCoreTests/EngineV2LivenessRecoveryTests.swift index 827e12685..d018bcd74 100644 --- a/provider-swift/Tests/ProviderCoreTests/EngineV2LivenessRecoveryTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/EngineV2LivenessRecoveryTests.swift @@ -350,8 +350,8 @@ struct EngineV2LivenessRecoveryTests { init() { // unloadModel / updateAggregateCapacity read MLX GPU counters — the - // mlx.metallib must be colocated with the test runner (CI extracts - // it; locally run `./scripts/fetch-metallib.sh debug` once). + // mlx.metallib must be colocated with the test runner (CI stages + // the source-built result; locally run `./scripts/fetch-metallib.sh debug` once). _ = LiveInferenceFixtures.ensureMetallibColocated() } diff --git a/provider-swift/Tests/ProviderCoreTests/EngineV2PagedParityLiveTests.swift b/provider-swift/Tests/ProviderCoreTests/EngineV2PagedParityLiveTests.swift index b5f6b7520..ec4e59927 100644 --- a/provider-swift/Tests/ProviderCoreTests/EngineV2PagedParityLiveTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/EngineV2PagedParityLiveTests.swift @@ -158,10 +158,10 @@ struct EngineV2PagedParityLiveTests { /// "Insufficient memory (X GB free, need Y GB) …" message. Everything /// else reaching this catch IS a load-path regression the loop gates /// exist to expose — explicit-paged refusal (the policy REFUSES instead - /// of degrading for an explicit `.paged` selection), VLM extraction or - /// engine-construction breakage, an invalid model directory, the - /// post-bridge headroom guard unloading a fresh paged slot — and must - /// fail the test, not return green. + /// of degrading for an explicit `.paged` selection), VLM serving-model + /// resolution or engine-construction breakage, an invalid model + /// directory, the post-bridge headroom guard unloading a fresh paged slot + /// — and must fail the test, not return green. private func triageLoopGateLoadFailure(_ error: Error, arm: String) { if case InferenceError.modelLoadFailed(let message) = error, message.hasPrefix("Insufficient memory (") @@ -521,9 +521,9 @@ struct EngineV2PagedParityLiveTests { try await loop.ensureModelLoaded(modelId: Self.gemmaModelID) } catch { // Same triage as the gpt-oss arm: only the pre-load free-memory - // refusal skips; explicit-paged refusal, VLM extraction failure, - // engine-construction breakage, an invalid model dir, and the - // post-bridge headroom guard all FAIL. See + // refusal skips; explicit-paged refusal, VLM serving-model + // resolution or engine-construction breakage, an invalid model + // dir, and the post-bridge headroom guard all FAIL. See // triageLoopGateLoadFailure. triageLoopGateLoadFailure(error, arm: "gemma-paged-loop") return diff --git a/provider-swift/Tests/ProviderCoreTests/EngineV2ProductionWiringTests.swift b/provider-swift/Tests/ProviderCoreTests/EngineV2ProductionWiringTests.swift index 057574759..49ebb609e 100644 --- a/provider-swift/Tests/ProviderCoreTests/EngineV2ProductionWiringTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/EngineV2ProductionWiringTests.swift @@ -426,40 +426,6 @@ struct EngineV2SlotBuildTests { #expect(events.first?.fields?["error_class"]?.description.contains("InitFailure") == true) } - @Test("VLM slot: extraction failure surfaces as vlm_extraction_failed refusal") - func vlmExtractionFailureRefusesLoudly() async throws { - let loop = try makeWiringLoop() - let runtime = EngineV2Runtime() - let telemetry = WiringTelemetrySink() - await loop.setEngineV2RuntimeForTesting(runtime) - await loop.setEngineV2SlotHooksForTesting( - ProviderLoop.EngineV2SlotHooks( - emitTelemetry: telemetry.callback(), - physicalMemoryBytes: wiringPhysicalBytes, - makeEngine: { _, _ in - // Stands in for any extraction failure (config decode, - // verify [.all] mismatch, forward-parity gate). - throw EngineV2VLMTextExtractionError.parityMismatch("scripted") - })) - - await #expect(throws: EngineV2VLMTextExtractionError.self) { - _ = try await loop.resliceAndBuildEngineV2SlotForTesting( - modelId: "gemma-4-26b-qat-4bit", - modelType: "gemma4", - isVLM: true, - container: makeStubContainer(), - tokenizer: TokenizerHandle(WiringStubTokenizer()), - sizing: makeSizing(weightsGiB: 15) - ) - } - #expect(await runtime.bridge(forModel: "gemma-4-26b-qat-4bit") == nil) - let refusal = telemetry.events.first { - $0.fields?["operation"]?.description == "engine_v2_refusal" - } - #expect(refusal != nil) - #expect(refusal?.severity == .error) - #expect(refusal?.fields?["reason"]?.description == "vlm_extraction_failed") - } @Test("production factory: unsupported model class throws (→ refusal)") func productionFactoryRejectsUnsupportedModel() { @@ -1834,9 +1800,9 @@ struct EngineV2RequestRoutingTests { /// These tests drive the REAL `updateAggregateCapacity` / `unloadModel` /// paths, which read MLX GPU counters — so the mlx.metallib must be -/// colocated with the test runner. CI places it under `.build` (see -/// ci.yml "Extract mlx.metallib"); locally run `./scripts/fetch-metallib.sh -/// debug` once. Mirrors the `LiveInferenceFixtures` pattern. +/// colocated with the test runner. CI uses the canonical source builder and +/// stages its result under `.build`; locally run +/// `./scripts/fetch-metallib.sh debug` once. Mirrors the `LiveInferenceFixtures` pattern. @Suite("EngineV2 production wiring: runtime guards", .serialized) struct EngineV2RuntimeGuardTests { diff --git a/provider-swift/Tests/ProviderCoreTests/EngineV2SSDPrefixCacheLiveTests.swift b/provider-swift/Tests/ProviderCoreTests/EngineV2SSDPrefixCacheLiveTests.swift index ff7e0612b..50193b598 100644 --- a/provider-swift/Tests/ProviderCoreTests/EngineV2SSDPrefixCacheLiveTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/EngineV2SSDPrefixCacheLiveTests.swift @@ -93,11 +93,9 @@ struct EngineV2SSDPrefixCacheLiveTests { static let gemmaQatModelID = "mlx-community/gemma-4-26B-A4B-it-qat-4bit" - /// Prod gemma-4 checkpoints are VLM builds: load through the VLM - /// factory and extract the CBv2-adapted text model over the SAME - /// weight arrays — the identical path the slot factory takes. Layer - /// kinds ALSO derived config-only (the production SSD construction - /// path for VLM slots) and pinned equal to engine truth. + /// Production Gemma 4 checkpoints are VLM builds. Their loaded wrapper + /// owns and exposes the exact CBv2-adapted text tower; SSD topology comes + /// from that same live module with no config re-decode or extraction. private func loadGemmaQat() async throws -> LiveModel { guard LiveInferenceFixtures.ensureMetallibColocated() != nil else { throw LiveFixtureSkip.missingMetallib @@ -118,20 +116,18 @@ struct EngineV2SSDPrefixCacheLiveTests { let tokenizer: TokenizerHandle = await container.perform { ctx in TokenizerHandle(ctx.tokenizer) } - let extraction = try EngineV2VLMTextExtraction.extractTextModel( - from: snapshot.model, modelDirectory: directory) - // The config-only derivation (what the slot factory hands the SSD - // cache for a VLM slot) must match engine truth. - let configKinds = EngineV2VLMTextExtraction.cbv2LayerKinds(modelDirectory: directory) - #expect(configKinds == extraction.model.cbv2LayerKinds, - "config-only layer kinds drifted from the extracted model's") + let wrapper = try #require(snapshot.model as? MLXVLM.Gemma4) + let textModel = wrapper.textModel + let direct = try EngineV2Factory.directServingModel( + model: wrapper, isVLM: true) + #expect(ObjectIdentifier(direct) == ObjectIdentifier(textModel)) return LiveModel( modelID: "gemma-4-26b-qat-4bit", container: container, - model: extraction.model, + model: textModel, tokenizer: tokenizer, eosTokenIds: snapshot.eosTokenIds, - layerKinds: extraction.model.cbv2LayerKinds, + layerKinds: textModel.cbv2LayerKinds, modelDirectory: directory) } diff --git a/provider-swift/Tests/ProviderCoreTests/FrozenReplayRealModelTests.swift b/provider-swift/Tests/ProviderCoreTests/FrozenReplayRealModelTests.swift index cba461a15..a39560ccb 100644 --- a/provider-swift/Tests/ProviderCoreTests/FrozenReplayRealModelTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/FrozenReplayRealModelTests.swift @@ -69,14 +69,16 @@ struct FrozenReplayRealModelTests { eosTokenIds: context.configuration.eosTokenIds, extraEOSTokens: context.configuration.extraEOSTokens.sorted()) } - let extraction = try EngineV2VLMTextExtraction.extractTextModel( - from: snapshot.model, - modelDirectory: directory) + let wrapper = try #require(snapshot.model as? MLXVLM.Gemma4) + let textModel = wrapper.textModel + let direct = try EngineV2Factory.directServingModel( + model: wrapper, isVLM: true) + #expect(ObjectIdentifier(direct) == ObjectIdentifier(textModel)) return Loaded( container: container, - model: extraction.model, - layerKinds: extraction.model.cbv2LayerKinds, - vocabularySize: extraction.model.vocabularySize) + model: textModel, + layerKinds: textModel.cbv2LayerKinds, + vocabularySize: textModel.vocabularySize) } private func makeBank(_ loaded: Loaded) -> CBv2LayerCacheBank { @@ -86,7 +88,7 @@ struct FrozenReplayRealModelTests { CBv2LayerCache(layerIndex: index, kind: kind) }) case let gemma as Gemma4TextModel: - return CBv2LayerCacheBank(caches: gemma.newCacheV2 { index, kind in + return CBv2LayerCacheBank(caches: try! gemma.newCacheV2 { index, kind in CBv2LayerCache(layerIndex: index, kind: kind) }) default: diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaMTPProductionLiveTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaMTPProductionLiveTests.swift index 94b350c51..309f7696d 100644 --- a/provider-swift/Tests/ProviderCoreTests/GemmaMTPProductionLiveTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/GemmaMTPProductionLiveTests.swift @@ -100,7 +100,6 @@ struct GemmaMTPProductionLiveTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: MTPProductionLiveFixtures.targetID, isVLM: true, - modelDirectory: targetDirectory, container: container, specDecPreparation: preparation, assistantLoader: AlwaysFailMTPAssistantLoader()) @@ -155,7 +154,7 @@ struct GemmaMTPProductionLiveTests { } @Test( - "VLM extraction and tool-templated decode retain greedy parity", + "shared VLM text tower and tool-templated decode retain greedy parity", .enabled( if: MTPProductionLiveFixtures.enabled, MTPProductionLiveFixtures.disabledReason)) diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaOptimizationConfigTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaOptimizationConfigTests.swift new file mode 100644 index 000000000..3ed7f6799 --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/GemmaOptimizationConfigTests.swift @@ -0,0 +1,240 @@ +import Darwin +import Testing +@testable import ProviderCore + +@Suite("Gemma optimization config") +struct GemmaOptimizationConfigTests { + @Test("missing optimization section enables the selected stack") + func missingSectionDefaultsOn() { + let config = ConfigManager.parse(""" + [provider] + name = "test-provider" + """) + + #expect(config.gemmaOptimizations == GemmaOptimizationSettings()) + #expect(config.gemmaOptimizations.prefillLayer18) + #expect(config.gemmaOptimizations.weightedR1) + } + + @Test("partial optimization section defaults each missing key on") + func partialSectionDefaultsMissingKeysOn() { + let layerOnly = ConfigManager.parse(""" + [provider] + name = "test-provider" + + [gemma_optimizations] + prefill_layer18 = false + """) + #expect(!layerOnly.gemmaOptimizations.prefillLayer18) + #expect(layerOnly.gemmaOptimizations.weightedR1) + + let weightedOnly = ConfigManager.parse(""" + [provider] + name = "test-provider" + + [gemma_optimizations] + weighted_r1 = false + """) + #expect(weightedOnly.gemmaOptimizations.prefillLayer18) + #expect(!weightedOnly.gemmaOptimizations.weightedR1) + } + + @Test("explicit optimization values are honored") + func explicitValues() { + let config = ConfigManager.parse(""" + [provider] + name = "test-provider" + + [gemma_optimizations] + prefill_layer18 = false + weighted_r1 = false + """) + + #expect(!config.gemmaOptimizations.prefillLayer18) + #expect(!config.gemmaOptimizations.weightedR1) + } + + @Test("optimization settings round trip with snake-case TOML keys") + func snakeCaseRoundTrip() { + let original = ProviderConfig( + provider: ProviderSettings(name: "test-provider"), + gemmaOptimizations: GemmaOptimizationSettings( + prefillLayer18: false, + weightedR1: true + ) + ) + + let toml = ConfigManager.serialize(original) + let decoded = ConfigManager.parse(toml) + + #expect(toml.contains("[gemma_optimizations]")) + #expect(toml.contains("prefill_layer18 = false")) + #expect(toml.contains("weighted_r1 = true")) + #expect(!toml.contains("gemmaOptimizations")) + #expect(!toml.contains("prefillLayer18")) + #expect(!toml.contains("weightedR1")) + #expect(decoded == original) + } +} + +@Suite("Gemma optimization environment") +struct GemmaOptimizationEnvironmentTests { + private let expectedKeys: Set = [ + "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL", + "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT", + "MLX_GATHER_QMM_EXPERT_SLICES", + ] + + @Test("projection emits exactly the three selected controls") + func exactProjection() { + let enabled = GemmaOptimizationEnvironment.projection( + for: GemmaOptimizationSettings() + ) + #expect(enabled == [ + "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL": "18", + "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT": "1", + "MLX_GATHER_QMM_EXPERT_SLICES": "1", + ]) + + let disabled = GemmaOptimizationEnvironment.projection( + for: GemmaOptimizationSettings( + prefillLayer18: false, + weightedR1: false + ) + ) + #expect(disabled == [ + "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL": "0", + "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT": "0", + "MLX_GATHER_QMM_EXPERT_SLICES": "0", + ]) + } + + @Test("weighted unsort and safe R1 are atomic in every projection") + func weightedR1IsAtomic() { + for enabled in [false, true] { + let projection = GemmaOptimizationEnvironment.projection( + for: GemmaOptimizationSettings(weightedR1: enabled) + ) + #expect( + projection[GemmaOptimizationEnvironment.weightedUnsortKey] + == projection[GemmaOptimizationEnvironment.safeR1Key] + ) + #expect(Set(projection.keys) == expectedKeys) + } + } + + @Test("apply overwrites every projected value") + func applyUsesOverwrite() throws { + var values: [String: String] = [:] + var overwrites: [String: Int32] = [:] + let settings = GemmaOptimizationSettings( + prefillLayer18: false, + weightedR1: true + ) + + try GemmaOptimizationEnvironment.apply(settings) { name, value, overwrite in + values[name] = value + overwrites[name] = overwrite + return 0 + } + + #expect(values == [ + "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL": "0", + "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT": "1", + "MLX_GATHER_QMM_EXPERT_SLICES": "1", + ]) + // The application boundary must hand the environment exactly what + // projection() reports, or the release matrix describes a dispatch + // that never happened. + #expect(values == GemmaOptimizationEnvironment.projection(for: settings)) + #expect(Set(overwrites.keys) == expectedKeys) + #expect(overwrites.values.allSatisfy { $0 == 1 }) + } + + @Test("a rejected key fails the whole application with its errno") + func applyRejectsPartialLatch() { + var attempted: [String: String] = [:] + let settings = GemmaOptimizationSettings() + + do { + try GemmaOptimizationEnvironment.apply(settings) { name, value, _ in + attempted[name] = value + return name == GemmaOptimizationEnvironment.safeR1Key ? ENOMEM : 0 + } + Issue.record("a rejected key must fail the whole application") + } catch let error as GemmaOptimizationEnvironment.ApplicationFailure { + #expect(error == GemmaOptimizationEnvironment.ApplicationFailure( + keys: [GemmaOptimizationEnvironment.safeR1Key], + code: ENOMEM + )) + } catch { + Issue.record("expected ApplicationFailure, got \(error)") + } + + // A rejected key never truncates the attempt, and the values offered + // stay exactly the projection. + #expect(attempted == GemmaOptimizationEnvironment.projection(for: settings)) + } + + @Test("every rejected key is reported, ordered, with the first errno") + func applyReportsAllRejectedKeys() { + var order: [String] = [] + + do { + try GemmaOptimizationEnvironment.apply(GemmaOptimizationSettings()) { + name, _, _ in + order.append(name) + switch name { + case GemmaOptimizationEnvironment.safeR1Key: return EINVAL + case GemmaOptimizationEnvironment.weightedUnsortKey: return EPERM + default: return 0 + } + } + Issue.record("rejected keys must fail the whole application") + } catch let error as GemmaOptimizationEnvironment.ApplicationFailure { + #expect(error == GemmaOptimizationEnvironment.ApplicationFailure( + keys: [ + GemmaOptimizationEnvironment.safeR1Key, + GemmaOptimizationEnvironment.weightedUnsortKey, + ], + code: EINVAL + )) + } catch { + Issue.record("expected ApplicationFailure, got \(error)") + } + + // Sorted application keeps the reported failure identical across runs + // despite per-process dictionary hash ordering. + #expect(order == expectedKeys.sorted()) + } + + @Test("failure description names the rejected keys and the errno") + func failureDescriptionIsPrecise() { + let failure = GemmaOptimizationEnvironment.ApplicationFailure( + keys: [ + GemmaOptimizationEnvironment.weightedUnsortKey, + GemmaOptimizationEnvironment.safeR1Key, + ], + code: ENOMEM + ) + + #expect(failure.description.contains( + GemmaOptimizationEnvironment.weightedUnsortKey)) + #expect(failure.description.contains( + GemmaOptimizationEnvironment.safeR1Key)) + #expect(failure.description.contains(String(cString: strerror(ENOMEM)))) + } + + @Test("projection excludes dropped packing and prefill controls") + func droppedControlsAreAbsent() { + let keys = Set(GemmaOptimizationEnvironment.projection( + for: GemmaOptimizationSettings() + ).keys) + + #expect(!keys.contains("MLX_GEMMA4_FUSED_EXPERT_GATE_UP")) + #expect(!keys.contains("MLX_GEMMA4_FUSED_DENSE_GATE_UP")) + #expect(!keys.contains("DARKBLOOM_GEMMA4_PREFILL_TAIL_ROWS")) + #expect(!keys.contains("DARKBLOOM_GEMMA4_PREFILL_LAST_QUERY")) + #expect(keys == expectedKeys) + } +} diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaOptimizationReportingTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaOptimizationReportingTests.swift new file mode 100644 index 000000000..858b257fd --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/GemmaOptimizationReportingTests.swift @@ -0,0 +1,126 @@ +import Foundation +import Testing +@testable import ProviderCore + +@Suite("Gemma optimization reporting") +struct GemmaOptimizationReportingTests { + @Test("reason precedence covers disabled, model eligibility, and AOT") + func reasonPrecedence() { + let report = GemmaOptimizationReport( + layer18Requested: false, + layer18Effective: false, + weightedUnsortRequested: true, + weightedUnsortEffective: false, + safeR1Requested: true, + safeR1GeometryEligible: true, + safeR1AOTAvailable: false, + safeR1NAXAvailable: true) + + #expect(report.layer18.reason == .disabled) + #expect(report.weightedUnsort.reason == .modelIneligible) + #expect(report.safeR1.reason == .aotUnavailable) + #expect(!report.layer18.effective) + #expect(!report.weightedUnsort.effective) + #expect(!report.safeR1.effective) + } + + @Test("NAX precedence and effective state render centrally") + func naxAndEffectiveRendering() { + let nax = GemmaOptimizationReport( + layer18Requested: true, + layer18Effective: true, + weightedUnsortRequested: true, + weightedUnsortEffective: true, + safeR1Requested: true, + safeR1GeometryEligible: true, + safeR1AOTAvailable: true, + safeR1NAXAvailable: true) + #expect(nax.safeR1.reason == .naxPrecedence) + #expect(!nax.safeR1.effective) + + let effective = GemmaOptimizationReport( + layer18Requested: true, + layer18Effective: true, + weightedUnsortRequested: true, + weightedUnsortEffective: true, + safeR1Requested: true, + safeR1GeometryEligible: true, + safeR1AOTAvailable: true, + safeR1NAXAvailable: false) + let rendered = effective.logLine(modelId: "gemma-4-production") + #expect(effective.states.allSatisfy { $0.effective }) + #expect(rendered.contains("layer18(requested=true,effective=true,reason=effective)")) + #expect(rendered.contains("weighted_unsort(requested=true,effective=true,reason=effective)")) + #expect(rendered.contains("safe_r1(requested=true,effective=true,reason=effective)")) + + let events = effective.telemetryEvents(modelId: "gemma-4-production") + #expect(events.count == 3) + #expect(events.allSatisfy { $0.kind == .engineHealth }) + #expect(events.allSatisfy { $0.fields?["model"]?.description == "gemma-4-production" }) + #expect(events.allSatisfy { $0.fields?["target"]?.description == "requested_1_effective_1" }) + #expect(events.allSatisfy { $0.fields?["reason"]?.description == "effective" }) + let allowedKeys = Set([ + "component", "operation", "backend", "model", "target", "reason", + ]) + #expect(events.allSatisfy { Set($0.fields?.keys.map { $0 } ?? []) == allowedKeys }) + } +} + +@Suite("Packaged retained-Gemma smoke") +struct PackagedRetainedGemmaSmokeTests { + private let expectedProjection = [ + GemmaOptimizationEnvironment.prefillLayer18Key: "18", + GemmaOptimizationEnvironment.weightedUnsortKey: "1", + GemmaOptimizationEnvironment.safeR1Key: "1", + ] + + @Test("synthetic retained config has an exact three-key projection") + func exactProjectionAndNoRejectedKeys() throws { + let config = try PackagedRuntimeSmoke.retainedConfiguration() + let decodedProjection = GemmaOptimizationEnvironment.projection( + for: config.gemmaOptimizations) + #expect(decodedProjection == expectedProjection) + try PackagedRuntimeSmoke.validateRetainedProjection(decodedProjection) + #expect( + PackagedRuntimeSmoke.rejectedEnvironmentKeys( + projection: expectedProjection, + environment: expectedProjection).isEmpty) + + var poisoned = expectedProjection + poisoned[GemmaOptimizationEnvironment.safeR1Key] = "poisoned" + #expect( + PackagedRuntimeSmoke.rejectedEnvironmentKeys( + projection: expectedProjection, + environment: poisoned) + == [GemmaOptimizationEnvironment.safeR1Key]) + } + + @Test("safe R1 gate requires requested packaged AOT and unarmed counters") + func safeR1Requirements() throws { + try PackagedRuntimeSmoke.validateSafeR1( + requested: true, aotAvailable: true, countersArmed: false) + + #expect(throws: PackagedRuntimeSmoke.VerificationError.safeR1NotRequested) { + try PackagedRuntimeSmoke.validateSafeR1( + requested: false, aotAvailable: true, countersArmed: false) + } + #expect(throws: PackagedRuntimeSmoke.VerificationError.safeR1AOTUnavailable) { + try PackagedRuntimeSmoke.validateSafeR1( + requested: true, aotAvailable: false, countersArmed: false) + } + #expect(throws: PackagedRuntimeSmoke.VerificationError.safeR1CountersArmed) { + try PackagedRuntimeSmoke.validateSafeR1( + requested: true, aotAvailable: true, countersArmed: true) + } + } + + @Test("signed-child marker is an exact output line") + func signedChildMarker() { + let exact = Data( + "noise\n\(PackagedRuntimeSmoke.gemmaOptimizationSuccessMarker)\npaged-kernel-runtime-smoke: ok\n".utf8) + let embedded = Data( + "prefix-\(PackagedRuntimeSmoke.gemmaOptimizationSuccessMarker)-suffix\n".utf8) + #expect(PackagedRuntimeSmoke.containsGemmaOptimizationSuccessMarker(exact)) + #expect(!PackagedRuntimeSmoke.containsGemmaOptimizationSuccessMarker(embedded)) + } +} diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaVLMEngineV2LiveTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaVLMEngineV2LiveTests.swift index ee3ac8556..8e688d6bb 100644 --- a/provider-swift/Tests/ProviderCoreTests/GemmaVLMEngineV2LiveTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/GemmaVLMEngineV2LiveTests.swift @@ -1,35 +1,17 @@ // Copyright © 2026 Eigen Labs. // -// Live (weight-gated) validation of the v0.7.2 Gemma 4 VLM → engine_v2 -// text-model extraction, on the EXACT checkpoints production serves -// (gemma-4-26b-qat-4bit / gemma-4-26b-8bit — both ship a vision tower, so -// the provider loads them through VLMModelFactory exactly like -// `ProviderLoop.loadModelContainer`). Two stages, mirroring the prod -// serving path (updated for the v0.7.5 ONE-ENGINE release): +// Live (weight-gated) validation of direct Gemma 4 VLM text-tower ownership +// on the exact qat-4bit / 8bit checkpoints production serves. Both load via +// `VLMModelFactory`, exactly like `ProviderLoop.loadModelContainer`. // -// (a) EXTRACTION + PARITY — extract the CBv2-adapted MLXLLM -// `Gemma4TextModel` over the wrapper's weight arrays; the extraction -// must not duplicate weights (MLX active memory grows by ~nothing), -// and the built-in load-time parity gate must pass (each side's greedy -// argmax sits in the other's top-5 at every probe position, max -// |Δlogit| bounded — see EngineV2VLMTextExtraction for why token-exact -// WRAPPER parity is structurally unattainable). +// (a) DIRECT OWNERSHIP + MEMORY — resolving the CBv2 serving model returns +// the wrapper's exact `textModel` object and does not construct or +// materialize a second tower. +// (b) V2 SERVE — that owned model serves a text request through the real +// production bridge/routing seam with deterministic greedy output. // -// (b) V2 SERVE — the extracted model must serve a text request through -// the REAL production seam (`EngineV2Bridge` + -// `MultiModelBatchSchedulerEngine.streamChatCompletion`, the exact -// engine construction the slot factory performs after its own -// extraction): non-empty greedy output, byte-identical across two -// identical submissions. -// -// DELETED with the legacy engine (v0.7.5 one-engine): the "v2 == legacy -// greedy" stage (its reference — a legacy `BatchScheduler` run over the -// same extracted module — no longer exists; the engine-repo v2-vs-legacy -// invariant it verified died with the legacy engine) and the legacy-vision -// interleave stage (media on a slot without a v2 bridge now throws the -// fail-loud "no serving engine for media" error instead of serving via the -// wrapper; vision-through-v2 interleave hygiene is pinned by -// GemmaVLMVisionEngineV2LiveTests). +// Vision-through-v2 image/video behavior and interleave hygiene remain pinned +// by GemmaVLMVisionEngineV2LiveTests and GemmaVLMVideoEngineV2LiveTests. // // Gated like the other multi-GB Gemma tests: DARKBLOOM_LIVE_MLX_TESTS + // DARKBLOOM_LIVE_MLX_GEMMA, and each checkpoint is skipped cleanly when not @@ -45,7 +27,7 @@ import Testing @testable import ProviderCore -@Suite("Gemma 4 VLM engine_v2 extraction (live)", .serialized) +@Suite("Gemma 4 VLM direct text tower (live)", .serialized) struct GemmaVLMEngineV2LiveTests { /// The two production Gemma 4 checkpoints (coordinator catalog ids @@ -57,19 +39,14 @@ struct GemmaVLMEngineV2LiveTests { private struct LoadedVLMSlot { let modelID: String - let directory: URL let container: ModelContainer let tokenizer: TokenizerHandle let model: any LanguageModel let eosTokenIds: Set } - /// Load a checkpoint EXACTLY the way the provider does for a VLM slot: - /// `VLMModelFactory` (the container the vision tower lives in). Loaded - /// BY HAND rather than through `LiveInferenceFixtures.loadBridge` - /// because stage (a) must run the extraction ITSELF, on the raw wrapper - /// handle, with clean before/after memory readings — the fixture's - /// production bridge would run the extraction internally first. + /// Load a checkpoint exactly as the provider does for a VLM slot: + /// `VLMModelFactory`, retaining the wrapper that owns both towers. private func loadVLMSlot(modelID: String, budgetBytes: Int) async throws -> LoadedVLMSlot { guard LiveInferenceFixtures.ensureMetallibColocated() != nil else { throw LiveFixtureSkip.missingMetallib @@ -95,7 +72,6 @@ struct GemmaVLMEngineV2LiveTests { } return LoadedVLMSlot( modelID: modelID, - directory: directory, container: container, tokenizer: tokenizer, model: snapshot.model, @@ -122,48 +98,36 @@ struct GemmaVLMEngineV2LiveTests { MLX.Memory.clearCache() } - /// Stage (a): weight-sharing extraction + the load-time parity gate. - /// - /// Two concerns, measured separately: - /// * WEIGHT-SHARING — extract with the parity gate OFF and require MLX - /// active memory to grow by no more than a small tolerance. Skeleton - /// construction is lazy; `update(parameters:)` re-points at the - /// wrapper's arrays; nothing multi-GiB may be retained (the removed - /// SwitchGLU fused gate+up cache was the v0.7.2 black hole — a - /// second weight copy would show up as ≥ the model size). - /// * PARITY GATE — extract again with the gate ON (production default) - /// and require it to pass and report a bounded max |Δlogit|. Both - /// extractions share the same wrapper arrays, so this is not a second - /// copy either; the returned model is the one stage (b) serves. - private func runExtractionStage( + /// Resolve the production serving model and prove it is the wrapper-owned + /// object. Merely exposing/resolving the tower must not add meaningful MLX + /// residency; a second checkpoint tower would exceed this allowance by + /// many GiB. + private func runDirectTowerStage( _ slot: LoadedVLMSlot - ) throws -> EngineV2VLMTextExtraction.Extraction { - // Weight-sharing invariant (parity gate off). + ) throws -> Gemma4TextModel { + let wrapper = try #require( + slot.model as? MLXVLM.Gemma4, + "production Gemma VLM checkpoint must load as MLXVLM.Gemma4") MLX.Memory.clearCache() let activeBefore = MLX.GPU.activeMemory - let noParity = try EngineV2VLMTextExtraction.extractTextModel( - from: slot.model, modelDirectory: slot.directory, - environment: ["DARKBLOOM_ENGINE_V2_VLM_PARITY_CHECK": "0"]) - #expect(noParity.parityMaxAbsLogitDiff == nil) + let owned = wrapper.textModel + let serving = try EngineV2Factory.directServingModel( + model: wrapper, isVLM: true) + let textModel = try #require(serving as? Gemma4TextModel) + + #expect(ObjectIdentifier(owned) == ObjectIdentifier(textModel)) + #expect(wrapper.textModel === owned) let growth = max(0, MLX.GPU.activeMemory - activeBefore) - let allowance = 1_536 * 1024 * 1024 + let allowance = 64 * 1024 * 1024 #expect( growth < allowance, Comment( - rawValue: "extraction grew MLX active memory by \(growth) bytes " - + "(> 1.5 GiB) — weights or a module cache were duplicated " - + "instead of shared")) - - // Parity gate (production default env). Returns the model stage - // (b) runs on. - let extraction = try EngineV2VLMTextExtraction.extractTextModel( - from: slot.model, modelDirectory: slot.directory) - let diff = try #require( - extraction.parityMaxAbsLogitDiff, "parity gate did not run under the default env") + rawValue: "direct tower resolution grew MLX active memory by " + + "\(growth) bytes (> 64 MiB), suggesting duplicate residency")) print( - "[gemma-vlm-v2] \(slot.modelID) parity max |Δlogit| = \(diff), " - + "weight-share growth = \(growth) bytes") - return extraction + "[gemma-vlm-v2] \(slot.modelID) direct tower identity passed; " + + "resolution growth = \(growth) bytes") + return textModel } /// Drive one OpenAI-shaped text request through the production routing @@ -208,15 +172,13 @@ struct GemmaVLMEngineV2LiveTests { return content } - /// Build the REAL production v2 engine + bridge over the extracted text - /// model — the same construction `EngineV2SlotFactory.makeProductionBridge` - /// performs right after ITS extraction, so stage (b) serves through the - /// exact one-engine seam production uses. + /// Build the real production v2 engine + bridge over the VLM-owned text + /// model, matching `EngineV2SlotFactory.makeProductionBridge`. private func makeBridge( - slot: LoadedVLMSlot, extracted: Gemma4TextModel + slot: LoadedVLMSlot, textModel: Gemma4TextModel ) throws -> EngineV2Bridge { let engine = try EngineV2Factory.makeProductionEngine( - model: extracted, + model: textModel, tokenizer: slot.tokenizer.inner, kvBytesCapacity: 4 * 1024 * 1024 * 1024, maxConcurrentRequests: Int(BackendSettings.defaultEngineV2MaxConcurrent) @@ -232,10 +194,10 @@ struct GemmaVLMEngineV2LiveTests { /// Structured bridge lifecycle: `shutdown()` (engine drain + pump /// teardown) awaited on every exit path. private func withBridge( - slot: LoadedVLMSlot, extracted: Gemma4TextModel, + slot: LoadedVLMSlot, textModel: Gemma4TextModel, _ body: (EngineV2Bridge) async throws -> Void ) async throws { - let bridge = try makeBridge(slot: slot, extracted: extracted) + let bridge = try makeBridge(slot: slot, textModel: textModel) do { try await body(bridge) } catch { @@ -245,22 +207,16 @@ struct GemmaVLMEngineV2LiveTests { await bridge.shutdown() } - /// Stages (a)+(b) for one checkpoint: measured extraction + parity - /// gate, then greedy serve determinism through the production seam. - private func runExtractionAndV2Serve(_ slot: LoadedVLMSlot) async throws { - // (a) extraction + load-time parity gate + weight-sharing invariant. - let extraction = try runExtractionStage(slot) + /// Direct identity/memory proof followed by deterministic production serve. + private func runDirectTowerAndV2Serve(_ slot: LoadedVLMSlot) async throws { + let textModel = try runDirectTowerStage(slot) - // (b) v2 serve: the extracted model, behind the real EngineV2 - // bridge, must produce non-empty greedy output through the - // production routing seam — and byte-identical output for two - // identical submissions (greedy decode is deterministic). - try await withBridge(slot: slot, extracted: extraction.model) { bridge in + try await withBridge(slot: slot, textModel: textModel) { bridge in let prompt = OpenAIMessageContent.text( "Count from one to five as digits separated by commas.") let v2Text = try await streamText( slot: slot, bridge: bridge, userContent: prompt, maxTokens: 32) - #expect(!v2Text.isEmpty, "v2 text serve over the extracted model produced no content") + #expect(!v2Text.isEmpty, "v2 text serve over the VLM-owned model produced no content") let v2TextAgain = try await streamText( slot: slot, bridge: bridge, userContent: prompt, maxTokens: 32) @@ -268,37 +224,37 @@ struct GemmaVLMEngineV2LiveTests { #expect( v2TextAgain == v2Text, Comment( - rawValue: "v2 greedy decode over the extracted model is " + rawValue: "v2 greedy decode over the VLM-owned model is " + "non-deterministic: \(v2Text.debugDescription) vs " + "\(v2TextAgain.debugDescription)")) } } - // MARK: - qat-4bit: extraction parity + v2 serve + // MARK: - qat-4bit: direct ownership + v2 serve @Test( - "qat-4bit: extraction parity and v2 serve determinism", + "qat-4bit: direct tower identity, memory, and v2 serve determinism", .enabled(if: LiveInferenceFixtures.gemmaTestsEnabled) ) - func qat4bitExtractionAndV2Serve() async throws { + func qat4bitDirectTowerAndV2Serve() async throws { try await withLoadedVLMSlot( modelID: Self.qat4bitModelID, budgetBytes: 48 * 1024 * 1024 * 1024 ) { slot in - try await runExtractionAndV2Serve(slot) + try await runDirectTowerAndV2Serve(slot) } } - // MARK: - 8bit: extraction parity + v2 serve + // MARK: - 8bit: direct ownership + v2 serve @Test( - "8bit: extraction parity and v2 serve determinism", + "8bit: direct tower identity, memory, and v2 serve determinism", .enabled(if: LiveInferenceFixtures.gemmaTestsEnabled) ) - func eightBitExtractionAndV2Serve() async throws { + func eightBitDirectTowerAndV2Serve() async throws { try await withLoadedVLMSlot( modelID: Self.eightBitModelID, budgetBytes: 64 * 1024 * 1024 * 1024 ) { slot in - try await runExtractionAndV2Serve(slot) + try await runDirectTowerAndV2Serve(slot) } } } diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaVLMParityProbeMemoryLiveTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaVLMParityProbeMemoryLiveTests.swift deleted file mode 100644 index d5223942f..000000000 --- a/provider-swift/Tests/ProviderCoreTests/GemmaVLMParityProbeMemoryLiveTests.swift +++ /dev/null @@ -1,281 +0,0 @@ -// Copyright © 2026 Eigen Labs. -// -// Live (weight-gated) regression for the v0.7.2 black-hole incident -// (7×64 GB gemma-4-26b-8bit + 1×36 GB qat-4bit boxes rejecting 100% of -// requests with the shared-KV capacity string from their first request, -// permanently). -// -// ROOT CAUSE (measured on the real 8-bit checkpoint, 2026-07-03): -// `MLXLMCommon.SwitchGLU` used to lazily concatenate a fused gate+up copy -// of its quantized expert weights on FIRST FORWARD and retain it on the -// module (~540 MB per MoE layer, ~15 GiB model-wide on gemma-4-26b-8bit). -// The v0.7.2 VLM text extraction created a SECOND SwitchGLU tree over the -// same weights, and the load-time parity probe ran one forward through -// EACH tree — materializing TWO ~15 GiB fused copies before the first -// request and pushing the box past the 90% unified-memory cap forever. -// -// The fused cache has since been DELETED from SwitchGLU (benchmarks showed -// ~0% decode win at its only active shape, B=1 solo decode, for 8–15 GiB -// of always-resident memory). The regression therefore tightens: the -// extraction + parity probe must retain (almost) NOTHING beyond the -// already-resident weights. This test asserts, on the real checkpoint: -// -// 1. GROWTH BOUND — extraction WITH the parity probe (production default) -// grows MLX active memory by no more than a small tolerance (rope -// tables / compiled-graph constants) — no multi-GiB retained state. -// 2. STEADY STATE AT LOAD — re-running both probe forwards afterwards -// grows active memory by ~nothing: no lazy build is waiting to fire -// mid-serving. -// 3. ADMISSION — a `GlobalKVCacheBudget` viewing the post-extraction MLX -// state through a simulated 64 GB profile admits a typical worst-case -// request reservation (the incident's first-request rejection, -// inverted). -// 4. CAPACITY CONSISTENCY — the weights-derived v2 static ceiling equals -// the shared gate's live headroom on the same profile: with no -// engine-retained overhead there is nothing to net out, and the -// heartbeat max and the gate agree by construction. -// -// Gated like the other multi-GB Gemma tests: DARKBLOOM_LIVE_MLX_TESTS + -// DARKBLOOM_LIVE_MLX_GEMMA; skipped cleanly when no checkpoint is cached. - -import Foundation -import MLX -import MLXLMCommon -import MLXNN -import MLXVLM -import Testing - -@testable import ProviderCore - -@Suite("Gemma 4 VLM parity-probe memory hygiene (live)", .serialized) -struct GemmaVLMParityProbeMemoryLiveTests { - - /// The incident checkpoint (catalog id `gemma-4-26b-8bit`); falls back to - /// the qat-4bit build when only that one is cached — the no-retained- - /// growth invariant is checkpoint-independent. - private static let preferredModelIDs = [ - LiveInferenceFixtures.gemmaModelID, - "mlx-community/gemma-4-26B-A4B-it-qat-4bit", - ] - - private struct LoadedVLMSlot { - let modelID: String - let directory: URL - let container: ModelContainer - let model: any LanguageModel - /// Scheduler-free sizing snapshot (v0.7.5 one-engine): the fp16 KV - /// rate + weight bytes production feeds the load gate, the re-slice, - /// and `makeEngineV2BridgeForSlot` (the legacy `BatchScheduler` - /// surfaces this test used to read are deleted). - let sizing: SlotSizingSnapshot - } - - private func loadFirstCachedVLMSlot() async throws -> LoadedVLMSlot { - guard LiveInferenceFixtures.ensureMetallibColocated() != nil else { - throw LiveFixtureSkip.missingMetallib - } - var located: (String, URL)? = nil - for modelID in Self.preferredModelIDs { - if case .found(let directory) = LiveInferenceFixtures.locate(modelID) { - located = (modelID, directory) - break - } - } - guard let (modelID, directory) = located else { - throw LiveFixtureSkip.modelNotInCache(Self.preferredModelIDs[0]) - } - #expect(ProviderLoop.modelIsVLM(at: directory)) - LiveInferenceFixtures.applyMemoryBudget(maxBytes: 64 * 1024 * 1024 * 1024) - - let container = try await VLMModelFactory.shared.loadContainer( - from: directory, using: LocalTokenizerLoader()) - // Same post-load sizing pass production runs (ProviderLoop+ - // ModelLoading): weight walk + engine-truth fp16 KV rate. Pure - // reads — nothing here may allocate MLX state, so it cannot - // perturb the growth baselines below. - let sizing = await SlotSizingSnapshot.build( - container: container, - modelPath: directory, - fallbackDefaultMaxTokens: 256) - let snapshot = await container.perform { ctx in - EngineV2ModelSnapshot( - model: ctx.model, - eosTokenIds: ctx.configuration.eosTokenIds, - extraEOSTokens: []) - } - return LoadedVLMSlot( - modelID: modelID, - directory: directory, - container: container, - model: snapshot.model, - sizing: sizing - ) - } - - private static func gib(_ bytes: Int) -> String { - String(format: "%.2f GiB", Double(bytes) / (1024 * 1024 * 1024)) - } - - @Test( - "extraction + parity probe leave the shared KV budget serveable (64 GB profile)", - .enabled(if: LiveInferenceFixtures.gemmaTestsEnabled) - ) - func parityProbeMemoryHygieneAndBudgetAdmission() async throws { - let slot = try await loadFirstCachedVLMSlot() - // No unload hop anymore: the multi-GB residency dies with `slot` at - // scope exit (the container is the only retainer); trim the probe's - // pool garbage on every exit path so the next serialized live suite - // starts from a clean pool. - defer { MLX.Memory.clearCache() } - try await runBody(slot) - } - - private func runBody(_ slot: LoadedVLMSlot) async throws { - // Mirror the production sequence: the load path's post-load headroom - // check trims the cold-load pool BEFORE the bridge build, so the - // pre-extraction baseline starts from a clean pool exactly like prod. - MLX.Stream().synchronize() - MLX.Memory.clearCache() - let activeBefore = MLX.GPU.activeMemory - let sysBefore = SystemMemory.availableBytes() ?? 0 - print( - "[parity-mem] pre-extraction: active=\(Self.gib(activeBefore)) " - + "cache=\(Self.gib(MLX.GPU.cacheMemory)) sysAvail=\(Self.gib(Int(sysBefore)))") - - // REAL extraction with the parity gate ON (production default env). - let extraction = try EngineV2VLMTextExtraction.extractTextModel( - from: slot.model, modelDirectory: slot.directory) - #expect(extraction.parityMaxAbsLogitDiff != nil) - - let activeAfter = MLX.GPU.activeMemory - let cacheAfter = MLX.GPU.cacheMemory - let activeGrowth = max(0, activeAfter - activeBefore) - print( - "[parity-mem] post-extraction: active=\(Self.gib(activeAfter)) " - + "cache=\(Self.gib(cacheAfter)) activeGrowth=\(Self.gib(activeGrowth))") - - // 1. GROWTH BOUND: the extraction shares the wrapper's weight arrays - // and retains no engine-side caches, so the only legitimate - // growth is small one-time state (rope tables, compiled-graph - // constants, probe stragglers). Pre-v0.7.3 this measured ~30 GiB - // (two fused copies); with the fused cache deleted outright the - // bar is a flat 2 GiB. - let tolerance = 2 * 1024 * 1024 * 1024 - #expect( - activeGrowth < tolerance, - Comment( - rawValue: "extraction + parity probe retained \(Self.gib(activeGrowth)) " - + "(> 2 GiB) — a weight or cache copy is being built " - + "(the v0.7.2 black-hole shape)")) - // The pool must also come back trimmed (post-probe clearCache). - #expect( - cacheAfter < 1 * 1024 * 1024 * 1024, - Comment( - rawValue: "probe left \(Self.gib(cacheAfter)) in the MLX pool — " - + "the post-probe clearCache is missing")) - - // 2. STEADY STATE AT LOAD: re-running both probe forwards must not - // materialize anything new — no lazy multi-GiB build can be - // waiting to fire on the first real request. - guard let wrapper = slot.model as? MLXVLM.Gemma4 else { - Issue.record("prod Gemma 4 slot did not load the MLXVLM.Gemma4 wrapper") - return - } - let probeTokens = MLXArray([2, 651, 6134, 1024, 578, 108, 2364].map(Int32.init)) - .expandedDimensions(axis: 0) - eval(wrapper(probeTokens, cache: nil)) - eval(extraction.model(probeTokens, cache: nil)) - MLX.Stream().synchronize() - MLX.Memory.clearCache() - let activeSteady = MLX.GPU.activeMemory - let steadyGrowth = max(0, activeSteady - activeAfter) - print("[parity-mem] steady-state re-forward growth=\(Self.gib(steadyGrowth))") - #expect( - steadyGrowth < 512 * 1024 * 1024, - Comment( - rawValue: "serving forwards after load grew active memory by " - + "\(Self.gib(steadyGrowth)) — a lazy per-tree cache still fires at " - + "request time")) - - // 3. ADMISSION on the incident profile: view the REAL post-extraction - // MLX counters through a synthetic 64 GB box (the incident - // hardware). The budget must admit one typical worst-case request - // (2 KiB prompt + 4096 max tokens at the slot's fp16 KV rate) — - // pre-fix this is exactly the reservation that failed 100% of the - // time from the first request. - let fp16Rate = slot.sizing.fp16KVBytesPerToken - let totalBytes: UInt64 = 64 * 1024 * 1024 * 1024 - let budget = GlobalKVCacheBudget( - memorySnapshot: { - let active = UInt64(max(0, MLX.GPU.activeMemory)) - let cache = UInt64(max(0, MLX.GPU.cacheMemory)) - let used = active + cache - // Simulated OS view of the 64 GB box: whatever the provider - // isn't holding, minus a 4 GiB OS/wired allowance. - let osAllowance: UInt64 = 4 * 1024 * 1024 * 1024 - let free = totalBytes > used + osAllowance ? totalBytes - used - osAllowance : 0 - return .init( - total: totalBytes, active: active, cache: cache, systemAvailable: free) - } - ) - let worstCaseTokens = 2048 + 4096 - let admitted = await budget.reserve( - requestID: "incident-probe", kvBytesPerToken: fp16Rate, tokenCount: worstCaseTokens) - print( - "[parity-mem] 64GB-profile admission: fp16KVBytesPerToken=\(fp16Rate) " - + "worstCaseTokens=\(worstCaseTokens) admitted=\(admitted)") - #expect( - admitted, - Comment( - rawValue: "the shared KV budget rejected a typical request on the simulated " - + "64 GB profile after the parity probe — the v0.7.2 black-hole signature")) - await budget.release(requestID: "incident-probe") - #expect(await budget.reservationIDsForTesting().isEmpty) - - // 4. CAPACITY CONSISTENCY: the v2 static ceiling is sized from the - // SIZING SNAPSHOT's weight figure (the exact input production - // makeEngineV2BridgeForSlot passes — ProviderLoop+EngineV2 reads - // `sizing.weightsBytes`), NOT from live MLX usage, so this - // genuinely catches the over-advertising shape: if - // extraction/probe retained any non-weight state, the - // weights-derived ceiling would exceed the gate's live headroom - // by exactly that retained amount. With the fused cache deleted, - // the divergence must fit inside the same 2 GiB retained-state - // bar as step 1 (rope tables, compiled-graph constants), and can - // never be negative beyond rounding (live use cannot be below - // the resident weights). - MLX.Stream().synchronize() - MLX.Memory.clearCache() - let snapshotWeightBytes = slot.sizing.weightsBytes - let mlxUsedNow = - UInt64(max(0, MLX.GPU.activeMemory)) + UInt64(max(0, MLX.GPU.cacheMemory)) - let ceiling = EngineV2KVSizing.engineKVBytesCapacity( - newModelWeightBytes: snapshotWeightBytes, - coResidentWeightBytes: 0, - existingEngineKVCapacities: [], - physicalBytes: totalBytes) - let gateHeadroom = UnifiedMemoryCap.liveKVHeadroomBytes( - physicalBytes: totalBytes, - mlxUsedBytes: mlxUsedNow, - systemAvailableBytes: .max) - let retainedOverhead = Int64(ceiling) - Int64(clamping: gateHeadroom) - print( - "[parity-mem] 64GB-profile capacity: ceiling=\(Self.gib(ceiling)) " - + "gateHeadroom=\(Self.gib(Int(gateHeadroom))) " - + "retainedOverhead=\(Self.gib(Int(retainedOverhead)))") - #expect( - retainedOverhead <= Int64(tolerance), - Comment( - rawValue: "weights-derived v2 ceiling \(Self.gib(ceiling)) exceeds the " - + "shared-gate live headroom \(Self.gib(Int(gateHeadroom))) by more than " - + "the retained-state bar — extraction/probe is holding non-weight " - + "memory the heartbeat would over-advertise (the finding-2 " - + "over-routing shape)")) - #expect( - retainedOverhead >= -(64 * 1024 * 1024), - Comment( - rawValue: "shared-gate headroom exceeds the weights-derived ceiling — " - + "live MLX usage measured below the snapshot's resident weights, " - + "which means the weight figure itself is inflated")) - } -} diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaVLMSharedTowerMemoryLiveTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaVLMSharedTowerMemoryLiveTests.swift new file mode 100644 index 000000000..4315b9817 --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/GemmaVLMSharedTowerMemoryLiveTests.swift @@ -0,0 +1,224 @@ +// Copyright © 2026 Eigen Labs. +// +// Live (weight-gated) regression for Gemma 4 VLM shared-tower residency. +// The loaded MLXVLM wrapper must expose its already-owned text tower to +// CBv2 without constructing a second module or retaining a second checkpoint +// worth of state. On a real production checkpoint this suite proves: +// +// 1. DIRECT RESOLUTION — production model resolution returns the wrapper's +// exact `textModel` object with negligible active-memory growth. +// 2. STEADY STATE — direct VLM and CBv2 forwards over that one object do not +// trigger a second lazy multi-GiB module cache. +// 3. ADMISSION — the post-forward shared KV budget remains serveable on the +// incident 64 GB profile. +// 4. CAPACITY CONSISTENCY — weights-derived capacity and live headroom stay +// aligned because no separately reconstructed tower is resident. +// +// Gated by DARKBLOOM_LIVE_MLX_TESTS + DARKBLOOM_LIVE_MLX_GEMMA and skipped +// cleanly when no checkpoint is cached. + +import Foundation +import MLX +import MLXLLM +import MLXLMCommon +import MLXNN +import MLXVLM +import Testing + +@testable import ProviderCore + +@Suite("Gemma 4 VLM shared-tower memory hygiene (live)", .serialized) +struct GemmaVLMSharedTowerMemoryLiveTests { + + /// The incident checkpoint (catalog id `gemma-4-26b-8bit`); falls back to + /// the qat-4bit build when only that one is cached — the no-retained- + /// growth invariant is checkpoint-independent. + private static let preferredModelIDs = [ + LiveInferenceFixtures.gemmaModelID, + "mlx-community/gemma-4-26B-A4B-it-qat-4bit", + ] + + private struct LoadedVLMSlot { + let modelID: String + let container: ModelContainer + let model: any LanguageModel + /// Scheduler-free sizing snapshot (v0.7.5 one-engine): the fp16 KV + /// rate + weight bytes production feeds the load gate, the re-slice, + /// and `makeEngineV2BridgeForSlot` (the legacy `BatchScheduler` + /// surfaces this test used to read are deleted). + let sizing: SlotSizingSnapshot + } + + private func loadFirstCachedVLMSlot() async throws -> LoadedVLMSlot { + guard LiveInferenceFixtures.ensureMetallibColocated() != nil else { + throw LiveFixtureSkip.missingMetallib + } + var located: (String, URL)? = nil + for modelID in Self.preferredModelIDs { + if case .found(let directory) = LiveInferenceFixtures.locate(modelID) { + located = (modelID, directory) + break + } + } + guard let (modelID, directory) = located else { + throw LiveFixtureSkip.modelNotInCache(Self.preferredModelIDs[0]) + } + #expect(ProviderLoop.modelIsVLM(at: directory)) + LiveInferenceFixtures.applyMemoryBudget(maxBytes: 64 * 1024 * 1024 * 1024) + + let container = try await VLMModelFactory.shared.loadContainer( + from: directory, using: LocalTokenizerLoader()) + // Same post-load sizing pass production runs (ProviderLoop+ + // ModelLoading): weight walk + engine-truth fp16 KV rate. Pure + // reads — nothing here may allocate MLX state, so it cannot + // perturb the growth baselines below. + let sizing = await SlotSizingSnapshot.build( + container: container, + modelPath: directory, + fallbackDefaultMaxTokens: 256) + let snapshot = await container.perform { ctx in + EngineV2ModelSnapshot( + model: ctx.model, + eosTokenIds: ctx.configuration.eosTokenIds, + extraEOSTokens: []) + } + return LoadedVLMSlot( + modelID: modelID, + container: container, + model: snapshot.model, + sizing: sizing + ) + } + + private static func gib(_ bytes: Int) -> String { + String(format: "%.2f GiB", Double(bytes) / (1024 * 1024 * 1024)) + } + + @Test( + "direct shared tower leaves the KV budget serveable (64 GB profile)", + .enabled(if: LiveInferenceFixtures.gemmaTestsEnabled) + ) + func sharedTowerMemoryHygieneAndBudgetAdmission() async throws { + let slot = try await loadFirstCachedVLMSlot() + // The multi-GB residency dies with `slot` at scope exit; trim forward + // buffers so the next serialized live suite starts from a clean pool. + defer { MLX.Memory.clearCache() } + try await runBody(slot) + } + + private func runBody(_ slot: LoadedVLMSlot) async throws { + // Mirror production: trim cold-load pool state before direct serving + // model resolution and take the residency baseline. + MLX.Stream().synchronize() + MLX.Memory.clearCache() + let activeBefore = MLX.GPU.activeMemory + let sysBefore = SystemMemory.availableBytes() ?? 0 + print( + "[shared-tower-mem] pre-resolution: active=\(Self.gib(activeBefore)) " + + "cache=\(Self.gib(MLX.GPU.cacheMemory)) sysAvail=\(Self.gib(Int(sysBefore)))") + + let wrapper = try #require( + slot.model as? MLXVLM.Gemma4, + "production Gemma 4 slot must load the MLXVLM.Gemma4 wrapper") + let owned = wrapper.textModel + let serving = try EngineV2Factory.directServingModel( + model: wrapper, isVLM: true) + let textModel = try #require(serving as? MLXLLM.Gemma4TextModel) + #expect(ObjectIdentifier(owned) == ObjectIdentifier(textModel)) + + let activeAfter = MLX.GPU.activeMemory + let activeGrowth = max(0, activeAfter - activeBefore) + print( + "[shared-tower-mem] post-resolution: active=\(Self.gib(activeAfter)) " + + "activeGrowth=\(Self.gib(activeGrowth))") + let tolerance = 64 * 1024 * 1024 + #expect( + activeGrowth < tolerance, + Comment( + rawValue: "direct tower resolution retained \(Self.gib(activeGrowth)) " + + "(> 64 MiB), suggesting a second module or weight copy")) + + // Direct VLM and CBv2 paths call the same object. A second forward must + // not materialize a second per-tree cache. + let probeTokens = MLXArray([2, 65, 61, 24, 57, 10, 23].map(Int32.init)) + .expandedDimensions(axis: 0) + eval(wrapper(probeTokens, cache: nil)) + eval(textModel(probeTokens, cache: nil)) + MLX.Stream().synchronize() + MLX.Memory.clearCache() + let activeSteady = MLX.GPU.activeMemory + let steadyGrowth = max(0, activeSteady - activeAfter) + print( + "[shared-tower-mem] steady-state re-forward growth=\(Self.gib(steadyGrowth))") + #expect( + steadyGrowth < 512 * 1024 * 1024, + Comment( + rawValue: "shared-tower forwards grew active memory by " + + "\(Self.gib(steadyGrowth)); a lazy duplicate cache fired")) + // Admission on the incident 64 GB profile: the live post-forward + // counters must admit a typical worst-case request. + let fp16Rate = slot.sizing.fp16KVBytesPerToken + let totalBytes: UInt64 = 64 * 1024 * 1024 * 1024 + let budget = GlobalKVCacheBudget( + memorySnapshot: { + let active = UInt64(max(0, MLX.GPU.activeMemory)) + let cache = UInt64(max(0, MLX.GPU.cacheMemory)) + let used = active + cache + // Simulated OS view of the 64 GB box: whatever the provider + // isn't holding, minus a 4 GiB OS/wired allowance. + let osAllowance: UInt64 = 4 * 1024 * 1024 * 1024 + let free = totalBytes > used + osAllowance ? totalBytes - used - osAllowance : 0 + return .init( + total: totalBytes, active: active, cache: cache, systemAvailable: free) + } + ) + let worstCaseTokens = 2048 + 4096 + let admitted = await budget.reserve( + requestID: "incident-probe", kvBytesPerToken: fp16Rate, tokenCount: worstCaseTokens) + print( + "[shared-tower-mem] 64GB-profile admission: fp16KVBytesPerToken=\(fp16Rate) " + + "worstCaseTokens=\(worstCaseTokens) admitted=\(admitted)") + #expect( + admitted, + Comment( + rawValue: "the shared KV budget rejected a typical request on the " + + "64 GB profile after direct shared-tower resolution")) + await budget.release(requestID: "incident-probe") + #expect(await budget.reservationIDsForTesting().isEmpty) + + // Capacity consistency: direct ownership leaves no separately + // reconstructed model state for the weights-derived ceiling to miss. + MLX.Stream().synchronize() + MLX.Memory.clearCache() + let snapshotWeightBytes = slot.sizing.weightsBytes + let mlxUsedNow = + UInt64(max(0, MLX.GPU.activeMemory)) + UInt64(max(0, MLX.GPU.cacheMemory)) + let ceiling = EngineV2KVSizing.engineKVBytesCapacity( + newModelWeightBytes: snapshotWeightBytes, + coResidentWeightBytes: 0, + existingEngineKVCapacities: [], + physicalBytes: totalBytes) + let gateHeadroom = UnifiedMemoryCap.liveKVHeadroomBytes( + physicalBytes: totalBytes, + mlxUsedBytes: mlxUsedNow, + systemAvailableBytes: .max) + let retainedOverhead = Int64(ceiling) - Int64(clamping: gateHeadroom) + print( + "[shared-tower-mem] 64GB-profile capacity: ceiling=\(Self.gib(ceiling)) " + + "gateHeadroom=\(Self.gib(Int(gateHeadroom))) " + + "retainedOverhead=\(Self.gib(Int(retainedOverhead)))") + let capacityTolerance = 2 * 1024 * 1024 * 1024 + #expect( + retainedOverhead <= Int64(capacityTolerance), + Comment( + rawValue: "weights-derived v2 ceiling \(Self.gib(ceiling)) exceeds the " + + "shared-gate live headroom \(Self.gib(Int(gateHeadroom))) by more " + + "than the retained-state bar, indicating unaccounted residency")) + #expect( + retainedOverhead >= -(64 * 1024 * 1024), + Comment( + rawValue: "shared-gate headroom exceeds the weights-derived ceiling — " + + "live MLX usage measured below the snapshot's resident weights, " + + "which means the weight figure itself is inflated")) + } +} diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaVLMVideoEngineV2LiveTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaVLMVideoEngineV2LiveTests.swift index 01efcac40..566d640a4 100644 --- a/provider-swift/Tests/ProviderCoreTests/GemmaVLMVideoEngineV2LiveTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/GemmaVLMVideoEngineV2LiveTests.swift @@ -4,10 +4,8 @@ // production Gemma 4 checkpoint (`gemma-4-26b-qat-4bit`). Self-contained // harness — deliberately independent of `GemmaVLMVisionEngineV2LiveTests` // (same precedent that file set vs the 0.7.3 suites; nothing here touches -// it). Slots are built through `LiveInferenceFixtures.loadBridge` — the -// production `EngineV2SlotFactory.makeProductionBridge` construction -// (weight-sharing text extraction + parity gate), i.e. the ONE engine -// every slot serves with since v0.7.5. Stages: +// it). Slots use `LiveInferenceFixtures.loadBridge`, which follows production +// direct shared-tower construction: one engine serves every slot request. // // (a) CONSTRUCTION INTROSPECTION — `EngineV2VisionPrefill.prepare` on a // real 4-distinct-frame clip: one span per sampled frame, spans == @@ -29,16 +27,10 @@ // measured and logged (informational — the vision tower runs // pre-submit). // -// DELETED with the legacy engine (v0.7.5 one-engine): the wrapper -// parity reference. Media on a slot WITHOUT a v2 bridge no longer -// serves at all — `MultiModelBatchSchedulerEngine` throws the -// fail-loud "no serving engine for media (no v2 bridge)" internal -// error (pinned by the non-live routing tests) — so there is no -// legacy output to log. (Token-exact wrapper parity was never -// asserted anyway: the extracted model implements the checkpoint's -// declared `rope_type: "proportional"` correctly while the wrapper -// deviated, plus bf16 kernel-order noise — see -// EngineV2VLMTextExtraction.) +// The removed legacy comparison is no longer needed for tower parity: +// direct VLM and CBv2 now invoke the same `Gemma4TextModel` instance. +// Media without a v2 bridge still fails loudly, as pinned by non-live +// routing tests. // // (c) 32-FRAME SAMPLING CAP — a 40-second, 40-frame clip: the processor // samples uniformly and caps at 32; construction must carve ≤ 32 diff --git a/provider-swift/Tests/ProviderCoreTests/GemmaVLMVisionEngineV2LiveTests.swift b/provider-swift/Tests/ProviderCoreTests/GemmaVLMVisionEngineV2LiveTests.swift index 042f4ee80..6918730b4 100644 --- a/provider-swift/Tests/ProviderCoreTests/GemmaVLMVisionEngineV2LiveTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/GemmaVLMVisionEngineV2LiveTests.swift @@ -25,15 +25,9 @@ // masks / spliced embeddings must leave no residue in the engine's // caches — the Qwen3.5-mrope-class regression pattern). // -// DELETED with the legacy engine (v0.7.5): the wrapper-path greedy -// comparison and the text/legacy-vision/text interleave. Media on a slot -// WITHOUT a v2 bridge no longer serves at all — it throws the fail-loud -// "no serving engine for media (no v2 bridge)" internal error, pinned by -// the non-live routing tests — so there is no legacy output to compare -// against. (Wrapper-vs-extracted token-exactness was never asserted anyway: -// the extracted model implements the checkpoint's declared `rope_type: -// "proportional"` correctly while the wrapper deviates, plus bf16 -// kernel-order noise — see EngineV2VLMTextExtraction.) +// The removed legacy comparison is no longer needed for tower parity: direct +// VLM and CBv2 now invoke the same `Gemma4TextModel` instance. Media without +// a v2 bridge still fails loudly, as pinned by non-live routing tests. // // Teardown here is structured (bridge shutdown awaited on every exit path) // so this suite never overlaps residency with other serialized live runs. @@ -82,10 +76,9 @@ struct GemmaVLMVisionEngineV2LiveTests { // MARK: - Harness (the provider's one-engine VLM slot shape) - /// One loaded v2 VLM slot: the PRODUCTION bridge (built through - /// `EngineV2SlotFactory.makeProductionBridge`, weight-sharing text - /// extraction + parity gate included), the retained VLM container the - /// vision tower runs in, and the tokenizer for the registry entry. + /// One loaded v2 VLM slot: the production bridge over the wrapper-owned + /// text tower, the retained VLM container that runs vision, and the + /// tokenizer for the registry entry. private struct LoadedV2VLMSlot { let modelID: String let bridge: EngineV2Bridge diff --git a/provider-swift/Tests/ProviderCoreTests/Helpers/MLXMetallibEnvironment.swift b/provider-swift/Tests/ProviderCoreTests/Helpers/MLXMetallibEnvironment.swift new file mode 100644 index 000000000..5d88a93fb --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/Helpers/MLXMetallibEnvironment.swift @@ -0,0 +1,45 @@ +import Foundation + +enum MLXMetallibEnvironment { + private static let key = "MLX_METALLIB_PATH" + private static let lock = NSRecursiveLock() + + static func withExclusiveAccess( + _ operation: () throws -> Result + ) rethrows -> Result { + lock.lock() + defer { lock.unlock() } + return try operation() + } + + static func setPath(_ path: String?) { + withExclusiveAccess { + updatePath(path) + } + } + + static func withPath( + _ path: String?, + operation: () throws -> Result + ) rethrows -> Result { + try withExclusiveAccess { + let previousPath = currentPath() + updatePath(path) + defer { updatePath(previousPath) } + return try operation() + } + } + + private static func currentPath() -> String? { + guard let value = getenv(key) else { return nil } + return String(cString: value) + } + + private static func updatePath(_ path: String?) { + if let path { + setenv(key, path, 1) + } else { + unsetenv(key) + } + } +} diff --git a/provider-swift/Tests/ProviderCoreTests/LaunchAgentRestartTests.swift b/provider-swift/Tests/ProviderCoreTests/LaunchAgentRestartTests.swift index e8a3f1988..d0fd26ae1 100644 --- a/provider-swift/Tests/ProviderCoreTests/LaunchAgentRestartTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/LaunchAgentRestartTests.swift @@ -1,3 +1,4 @@ +import Foundation import Testing @testable import ProviderCore @@ -82,6 +83,16 @@ struct LaunchAgentEnvironmentTests { #expect(out == ["DARKBLOOM_MTP_MAX_RECTANGULAR_TOKENS": "4"]) } + @Test func excludesConfigBackedGemmaControlsFromDaemonEnvironment() { + let out = LaunchAgent.passthroughEnvironment(from: [ + "DARKBLOOM_PREFIX_CACHE": "0", + GemmaOptimizationEnvironment.prefillLayer18Key: "poison", + GemmaOptimizationEnvironment.weightedUnsortKey: "poison", + GemmaOptimizationEnvironment.safeR1Key: "poison", + ]) + #expect(out == ["DARKBLOOM_PREFIX_CACHE": "0"]) + } + @Test func forwardsKVBackendGuardPathToDaemonAndWatchdog() { // The crash-loop guard record has one writer (the launchd watchdog) // and several readers (the launchd daemon's engine factory, a @@ -103,7 +114,13 @@ struct LaunchAgentServicePlistTests { label: "io.darkbloom.provider", programArguments: ["/usr/local/bin/darkbloom", "start", "--foreground"], logPath: "/tmp/p.log", - environment: ["DARKBLOOM_PREFIX_CACHE": "0", "PATH": "/usr/bin"] + environment: [ + "DARKBLOOM_PREFIX_CACHE": "0", + GemmaOptimizationEnvironment.prefillLayer18Key: "poison", + GemmaOptimizationEnvironment.weightedUnsortKey: "poison", + GemmaOptimizationEnvironment.safeR1Key: "poison", + "PATH": "/usr/bin", + ] ) // RunAtLoad=true so a rebooted / auto-login box restarts (and re-attests via // APNs) with no human; KeepAlive stays false to avoid racing the self-updater. @@ -122,4 +139,27 @@ struct LaunchAgentServicePlistTests { #expect(plist["EnvironmentVariables"] == nil) #expect(plist["RunAtLoad"] as? Bool == true) } + + @Test func customConfigFlagAndAbsolutePathStayAdjacent() throws { + let arguments = LaunchAgent.serviceProgramArguments( + binaryPath: "/usr/local/bin/darkbloom", + coordinatorURL: "wss://api.darkbloom.dev/ws/provider", + models: ["org/model"], + idleTimeout: 15, + configPath: URL(fileURLWithPath: "/tmp/custom provider.toml") + ) + let flagIndex = try #require(arguments.firstIndex(of: "--config")) + #expect(arguments[flagIndex + 1] == "/tmp/custom provider.toml") + } + + @Test func defaultConfigPathRemainsImplicit() { + let arguments = LaunchAgent.serviceProgramArguments( + binaryPath: "/usr/local/bin/darkbloom", + coordinatorURL: "wss://api.darkbloom.dev/ws/provider", + models: [], + idleTimeout: nil, + configPath: nil + ) + #expect(!arguments.contains("--config")) + } } diff --git a/provider-swift/Tests/ProviderCoreTests/LiveInferenceFixtures.swift b/provider-swift/Tests/ProviderCoreTests/LiveInferenceFixtures.swift index 3450b7ce9..75a2c254f 100644 --- a/provider-swift/Tests/ProviderCoreTests/LiveInferenceFixtures.swift +++ b/provider-swift/Tests/ProviderCoreTests/LiveInferenceFixtures.swift @@ -7,7 +7,6 @@ import MLXLMCommon // MARK: - Tiny, fast model used for the bulk of live tests. enum LiveInferenceFixtures { - /// Default tiny MLX-community model: ~600M params, ~1 GB on disk in 8-bit. /// Loads in seconds and finishes a 16-token generation in well under 1s /// on Apple Silicon. Has a chat template; no tool-calling weirdness. @@ -78,51 +77,55 @@ enum LiveInferenceFixtures { /// /// `.build//debug/PackageTests.xctest/Contents/MacOS/` /// - /// `scripts/fetch-metallib.sh` only places the metallib at - /// `.build/debug/mlx.metallib`, which is *not* where MLX looks. This - /// helper finds the metallib in any well-known location (incl. the - /// fetch script's drop site) and copies it next to the test runner - /// so MLX's `current_binary_dir() / "mlx.metallib"` lookup succeeds - /// on the first GPU call. Idempotent. + /// The canonical source helper stages a metallib in the Swift build + /// directory. This helper finds that source and always replaces the copy + /// beside the test runner so a stale pre-existing file cannot survive into + /// the current test invocation. /// - /// Returns the path to the colocated metallib on success, or `nil` if - /// no source metallib could be found anywhere -- in which case the - /// caller should skip the test rather than crashing in the GPU init. + /// Returns the path to the colocated metallib on success, or `nil` if no + /// source metallib could be found -- in which case the caller should skip + /// the test rather than crashing in GPU initialization. static func ensureMetallibColocated() -> URL? { - let fm = FileManager.default - - // 1. Find the test bundle's MacOS dir. Bundle(for:) reliably points - // at the .xctest bundle even when launched via the system - // `xctest` host (where _NSGetExecutablePath returns the host). - guard let testBundleMacOSDir = testBundleExecutableDir() else { - return nil - } - let destination = testBundleMacOSDir.appendingPathComponent("mlx.metallib") - - if fm.fileExists(atPath: destination.path) { - return destination - } + MLXMetallibEnvironment.withExclusiveAccess { + let fm = FileManager.default + + // 1. Find the test bundle's MacOS dir. Bundle(for:) reliably points + // at the .xctest bundle even when launched via the system + // `xctest` host (where _NSGetExecutablePath returns the host). + guard let testBundleMacOSDir = testBundleExecutableDir() else { + return nil + } + let destination = testBundleMacOSDir.appendingPathComponent("mlx.metallib") - // 2. Find a source metallib to copy. - guard let source = findSourceMetallib() else { - return nil - } + // 2. Resolve the authoritative staged source before considering the + // runner copy. Existence alone does not prove source compatibility. + guard let source = findSourceMetallib() else { + return nil + } - do { - try fm.copyItem(at: source, to: destination) - // Mirror to MLX_METALLIB_PATH so our own `locateMetallib()` - // (which trusts _NSGetExecutablePath, i.e. the xctest host - // path) can find it too if anyone else queries. - setenv("MLX_METALLIB_PATH", destination.path, 1) - return destination - } catch { - // Last resort: still set MLX_METALLIB_PATH so any code that - // queries `locateMetallib()` succeeds, even though the MLX - // C++ runtime won't honor it and tests will crash on first - // GPU call. Better to let the test report the failure than - // to silently skip. - setenv("MLX_METALLIB_PATH", source.path, 1) - return nil + do { + // Copy beside the destination, then atomically replace the runner + // file without loading the 150 MB+ metallib into process memory. + let temporary = testBundleMacOSDir + .appendingPathComponent(".mlx.metallib.\(UUID().uuidString)") + defer { try? fm.removeItem(at: temporary) } + try fm.copyItem(at: source, to: temporary) + if fm.fileExists(atPath: destination.path) { + _ = try fm.replaceItemAt(destination, withItemAt: temporary) + } else { + try fm.moveItem(at: temporary, to: destination) + } + // Mirror to MLX_METALLIB_PATH so our own `locateMetallib()` + // (which trusts _NSGetExecutablePath, i.e. the xctest host + // path) can find it too if anyone else queries. + MLXMetallibEnvironment.setPath(destination.path) + return destination + } catch { + // The C++ runtime does not honor MLX_METALLIB_PATH, so failure to + // replace its runner-local copy must remain a fixture failure. + MLXMetallibEnvironment.setPath(source.path) + return nil + } } } @@ -146,29 +149,33 @@ enum LiveInferenceFixtures { return nil } - /// Look for a metallib at the spots `scripts/fetch-metallib.sh` and the - /// release pipeline drop one. We anchor at the test bundle (which is - /// inside `.build//debug/...`) and walk up to the package root. + /// Look for a metallib at the canonical helper's drop sites. We anchor at + /// the test bundle (`.build///...`) and accept only + /// the configuration which contains the running test bundle. private static func findSourceMetallib() -> URL? { let fm = FileManager.default - if let env = ProcessInfo.processInfo.environment["MLX_METALLIB_SOURCE"], - !env.isEmpty, - fm.fileExists(atPath: env) { - return URL(fileURLWithPath: env) - } - // Anchor at the test bundle path -- much more reliable than // _NSGetExecutablePath under `swift test`. let bundle = Bundle(for: BundleSentinel.self) + let components = bundle.bundleURL.pathComponents + let configuration: String + if let buildIndex = components.lastIndex(of: ".build"), + let activeConfiguration = components[components.index(after: buildIndex)...] + .first(where: { $0 == "debug" || $0 == "release" }) { + configuration = activeConfiguration + } else { + configuration = "debug" + } + var cursor = bundle.bundleURL for _ in 0..<12 { if cursor.lastPathComponent == ".build" { let candidates: [URL] = [ - cursor.appendingPathComponent("debug/mlx.metallib"), - cursor.appendingPathComponent("release/mlx.metallib"), - cursor.appendingPathComponent("arm64-apple-macosx/debug/mlx.metallib"), - cursor.appendingPathComponent("arm64-apple-macosx/release/mlx.metallib"), + cursor.appendingPathComponent("\(configuration)/mlx.metallib"), + cursor.appendingPathComponent( + "arm64-apple-macosx/\(configuration)/mlx.metallib" + ), ] for candidate in candidates { if fm.fileExists(atPath: candidate.path) { @@ -215,9 +222,8 @@ enum LiveInferenceFixtures { let sizing: SlotSizingSnapshot } - /// Load a model and build its production v2 bridge (VLM-aware — the - /// same `ModelContainerLoading` + weight-sharing text extraction the - /// serve path uses). + /// Load a model and build its production v2 bridge, including direct use + /// of a VLM wrapper's owned text tower. /// /// - Throws: `LiveFixtureSkip` if the model isn't on disk, or if the /// metallib isn't available. diff --git a/provider-swift/Tests/ProviderCoreTests/MTPConfigTests.swift b/provider-swift/Tests/ProviderCoreTests/MTPConfigTests.swift index ed935d4e4..a85f3884d 100644 --- a/provider-swift/Tests/ProviderCoreTests/MTPConfigTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/MTPConfigTests.swift @@ -174,12 +174,16 @@ struct MTPBetaFeatureTests { feature.apply(true, to: &config) #expect(config.backend.mtp == true) #expect(feature.isEnabled(in: config) == true) - #expect(BetaFeatures.enabledIDs(in: config) == ["mtp"]) + #expect(BetaFeatures.enabledIDs(in: config) == [ + "gemma-prefill-layer18", "gemma-weighted-r1", "mtp", + ]) feature.apply(false, to: &config) #expect(config.backend.mtp == false) #expect(feature.isEnabled(in: config) == false) - #expect(BetaFeatures.enabledIDs(in: config).isEmpty) + #expect(BetaFeatures.enabledIDs(in: config) == [ + "gemma-prefill-layer18", "gemma-weighted-r1", + ]) } @Test("apply only mutates its mapped field") @@ -195,6 +199,7 @@ struct MTPBetaFeatureTests { #expect(config.backend.port == before.backend.port) #expect(config.provider == before.provider) #expect(config.coordinator == before.coordinator) + #expect(config.gemmaOptimizations == before.gemmaOptimizations) } // `darkbloom beta enable mtp` = apply(true) + ConfigManager.save; the diff --git a/provider-swift/Tests/ProviderCoreTests/MetallibHashTests.swift b/provider-swift/Tests/ProviderCoreTests/MetallibHashTests.swift index a6a46df34..d5e575536 100644 --- a/provider-swift/Tests/ProviderCoreTests/MetallibHashTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/MetallibHashTests.swift @@ -2,10 +2,8 @@ import Foundation import Testing @testable import ProviderCore -/// Marked `.serialized` because every test in this suite mutates the -/// process-wide MLX_METALLIB_PATH environment variable. Swift Testing's -/// default parallel execution would race them and produce flakes like -/// "metallibHash returned nil because another test just unset the env". +/// Marked `.serialized` to avoid needless contention within this suite. The +/// shared environment guard also excludes mutations from live-test fixtures. @Suite("metallib hash + locator", .serialized) struct MetallibHashTests { @@ -16,11 +14,10 @@ struct MetallibHashTests { defer { try? FileManager.default.removeItem(at: tmp) } try Data("not really a metallib but exists".utf8).write(to: tmp) - setenv("MLX_METALLIB_PATH", tmp.path, 1) - defer { unsetenv("MLX_METALLIB_PATH") } - - let located = locateMetallib() - #expect(located?.path == tmp.path) + MLXMetallibEnvironment.withPath(tmp.path) { + let located = locateMetallib() + #expect(located?.path == tmp.path) + } } @Test("metallibHash returns a 64-character hex string when located") @@ -30,16 +27,15 @@ struct MetallibHashTests { defer { try? FileManager.default.removeItem(at: tmp) } try Data(repeating: 0x42, count: 1024).write(to: tmp) - setenv("MLX_METALLIB_PATH", tmp.path, 1) - defer { unsetenv("MLX_METALLIB_PATH") } - - guard let hash = metallibHash() else { - Issue.record("metallibHash returned nil for an existing file at \(tmp.path)") - return + MLXMetallibEnvironment.withPath(tmp.path) { + guard let hash = metallibHash() else { + Issue.record("metallibHash returned nil for an existing file at \(tmp.path)") + return + } + #expect(hash.count == 64) + let hex = Set("0123456789abcdef") + #expect(hash.allSatisfy { hex.contains($0) }) } - #expect(hash.count == 64) - let hex = Set("0123456789abcdef") - #expect(hash.allSatisfy { hex.contains($0) }) } @Test("metallibHash is stable across calls for the same file") @@ -49,13 +45,32 @@ struct MetallibHashTests { defer { try? FileManager.default.removeItem(at: tmp) } try Data("hello mlx".utf8).write(to: tmp) - setenv("MLX_METALLIB_PATH", tmp.path, 1) - defer { unsetenv("MLX_METALLIB_PATH") } + let competing = URL(fileURLWithPath: NSTemporaryDirectory()) + .appendingPathComponent("competing-mlx-\(UUID().uuidString).metallib") + defer { try? FileManager.default.removeItem(at: competing) } + try Data("different metallib".utf8).write(to: competing) + + let mutationAttempted = DispatchSemaphore(value: 0) + let mutationCompleted = DispatchSemaphore(value: 0) + + MLXMetallibEnvironment.withPath(tmp.path) { + let firstHash = metallibHash() + let mutationThread = Thread { + mutationAttempted.signal() + MLXMetallibEnvironment.withPath(competing.path) {} + mutationCompleted.signal() + } + mutationThread.start() - let a = metallibHash() - let b = metallibHash() - #expect(a != nil) - #expect(a == b) + #expect(mutationAttempted.wait(timeout: .now() + 5) == .success) + #expect(mutationCompleted.wait(timeout: .now() + 0.05) == .timedOut) + + let secondHash = metallibHash() + #expect(firstHash != nil) + #expect(firstHash == secondHash) + } + + #expect(mutationCompleted.wait(timeout: .now() + 5) == .success) } @Test("locateMetallib returns nil when nothing is found and no env override") @@ -64,13 +79,12 @@ struct MetallibHashTests { // through to the binary-adjacent search and may or may not find one // (it could find one in the test bundle's .build path). We assert // on the env override semantics only. - setenv("MLX_METALLIB_PATH", "/var/empty/definitely-not-here.metallib", 1) - defer { unsetenv("MLX_METALLIB_PATH") } - - // Env override misses → falls back to binary-adjacent search. The - // test binary may or may not have a colocated metallib; we don't - // assert one way or the other, just that the function returns - // without crashing. - _ = locateMetallib() + MLXMetallibEnvironment.withPath("/var/empty/definitely-not-here.metallib") { + // Env override misses → falls back to binary-adjacent search. The + // test binary may or may not have a colocated metallib; we don't + // assert one way or the other, just that the function returns + // without crashing. + _ = locateMetallib() + } } } diff --git a/provider-swift/Tests/ProviderCoreTests/PagedDivergenceProbeTests.swift b/provider-swift/Tests/ProviderCoreTests/PagedDivergenceProbeTests.swift index 9c885d13a..169b902eb 100644 --- a/provider-swift/Tests/ProviderCoreTests/PagedDivergenceProbeTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/PagedDivergenceProbeTests.swift @@ -293,7 +293,7 @@ struct PagedDivergenceProbeTests { } let box = try await container.perform { ctx -> Box in let serving = try EngineV2Factory.benchmarkServingModel( - model: ctx.model, isVLM: isVLM, modelDirectory: directory) + model: ctx.model, isVLM: isVLM) guard let kinds = EngineV2Factory.cbv2LayerKinds(model: serving) else { throw LiveFixtureSkip.modelNotInCache("no cbv2 layer kinds for \(modelID)") } diff --git a/provider-swift/Tests/ProviderCoreTests/ProviderMTPFactoryTests.swift b/provider-swift/Tests/ProviderCoreTests/ProviderMTPFactoryTests.swift index 6dfdd2685..811913a77 100644 --- a/provider-swift/Tests/ProviderCoreTests/ProviderMTPFactoryTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/ProviderMTPFactoryTests.swift @@ -1,8 +1,10 @@ import Foundation import Crypto import MLX +import MLXLLM import MLXLMCommon import MLXNN +import MLXVLM import Testing @testable import ProviderCore import ProviderCoreFoundation @@ -33,7 +35,9 @@ private struct MTPFactoryProcessor: UserInputProcessor { func prepare(input: UserInput) async throws -> LMInput { throw CancellationError() } } -private func mtpFactoryContainer(_ target: MTPFactoryTarget = MTPFactoryTarget()) -> ModelContainer { +private func mtpFactoryContainer( + _ target: any LanguageModel = MTPFactoryTarget() +) -> ModelContainer { ModelContainer(context: ModelContext( configuration: ModelConfiguration(id: "test/mtp-factory"), model: target, @@ -41,6 +45,48 @@ private func mtpFactoryContainer(_ target: MTPFactoryTarget = MTPFactoryTarget() tokenizer: MTPFactoryTokenizer())) } +private func mtpFactoryVLM() throws -> MLXVLM.Gemma4 { + let data = Data( + """ + { + "model_type": "gemma4", + "text_config": { + "hidden_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "global_head_dim": 8, + "intermediate_size": 64, + "vocab_size": 128, + "sliding_window": 16, + "layer_types": ["sliding_attention", "full_attention"], + "tie_word_embeddings": true, + "hidden_size_per_layer_input": 0, + "vocab_size_per_layer_input": 0, + "num_kv_shared_layers": 0, + "use_double_wide_mlp": false, + "enable_moe_block": false, + "use_bidirectional_attention": "vision" + }, + "vision_config": { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "head_dim": 8, + "patch_size": 8, + "position_embedding_size": 64, + "default_output_length": 4, + "pooling_kernel_size": 2 + } + } + """.utf8) + return MLXVLM.Gemma4( + try JSONDecoder().decode(MLXVLM.Gemma4Configuration.self, from: data)) +} + private func mtpFactoryArtifact() throws -> SpecDecArtifact { let directory = FileManager.default.temporaryDirectory .appendingPathComponent("provider-mtp-test-\(UUID().uuidString)", isDirectory: true) @@ -207,7 +253,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: model.id, isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: preparation, assistantLoader: MTPFactoryRecordingLoader(recorder: recorder, failure: nil)) @@ -227,7 +272,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(target), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact)), @@ -240,6 +284,51 @@ struct ProviderMTPFactoryTests { #expect(prepared.assistantBytes == artifact.residentBytes) } + @Test("Gemma 4 VLM hands its owned text tower directly to CBv2 and MTP") + func vlmOwnedTargetIdentity() async throws { + let vlm = try mtpFactoryVLM() + let container = mtpFactoryContainer(vlm) + let recorder = MTPFactoryIdentityRecorder() + let artifact = try mtpFactoryArtifact() + defer { try? FileManager.default.removeItem(at: artifact.directory) } + + let prepared = try await EngineV2SlotFactory.prepareProductionModel( + modelId: "gemma-4-vlm-test", + isVLM: true, + container: container, + specDecPreparation: .init( + artifact: artifact, status: .candidate(artifact)), + assistantLoader: MTPFactoryRecordingLoader( + recorder: recorder, failure: nil)) + + let ownedID = ObjectIdentifier(vlm.textModel) + #expect(ObjectIdentifier(prepared.servingModel) == ownedID) + #expect(recorder.snapshot.0 == ownedID) + #expect(recorder.snapshot.1 == 1) + #expect(prepared.mtpStatus.active) + + let recovered = try await EngineV2SlotFactory.prepareRecoveryModel( + modelId: "gemma-4-vlm-test", + isVLM: true, + container: container, + previousArtifact: prepared.mtpArtifact, + previousStatus: prepared.mtpStatus, + assistant: prepared.assistant) + #expect(ObjectIdentifier(recovered.servingModel) == ownedID) + let recoveredAssistant = try #require(recovered.assistant) + let preparedAssistant = try #require(prepared.assistant) + #expect(recoveredAssistant === preparedAssistant) + #expect(recovered.mtpStatus.active) + } + + @Test("VLM resolution fails loud for an unsupported wrapper") + func unsupportedVLMWrapperRefuses() { + #expect(throws: EngineV2ProductionError.self) { + _ = try EngineV2Factory.directServingModel( + model: MTPFactoryTarget(), isVLM: true) + } + } + @Test("assistant load and bind failures are stable target-only fallbacks") func loadAndBindFallbacks() async throws { for failure in [ @@ -252,7 +341,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact)), @@ -272,7 +360,6 @@ struct ProviderMTPFactoryTests { let first = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact)), @@ -291,7 +378,6 @@ struct ProviderMTPFactoryTests { let rebuilt = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact)), @@ -317,7 +403,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact)), @@ -342,7 +427,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact)), @@ -361,7 +445,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: artifact, status: .candidate(artifact))) @@ -375,7 +458,6 @@ struct ProviderMTPFactoryTests { let prepared = try await EngineV2SlotFactory.prepareProductionModel( modelId: "gemma-4-test", isVLM: false, - modelDirectory: nil, container: mtpFactoryContainer(), specDecPreparation: .init( artifact: nil, diff --git a/provider-swift/Tests/ProviderCoreTests/SelfUpdaterTests.swift b/provider-swift/Tests/ProviderCoreTests/SelfUpdaterTests.swift index b32363143..2942ad027 100644 --- a/provider-swift/Tests/ProviderCoreTests/SelfUpdaterTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/SelfUpdaterTests.swift @@ -514,7 +514,7 @@ struct SelfUpdaterTests { return (tarball, release, install) } - @Test("signed extracted app runs real packaged verification before staging succeeds") + @Test("signed extracted child proves retained Gemma marker before staging succeeds") func signedAppRunsRealVerification() throws { _ = LiveInferenceFixtures.ensureMetallibColocated() let root = FileManager.default.temporaryDirectory @@ -532,7 +532,7 @@ struct SelfUpdaterTests { release: validRelease, installDir: install) else { - Issue.record("real signed/runtime-verified staging failed") + Issue.record("real signed/runtime marker verification failed") return } staged.discard() diff --git a/provider-swift/Tests/ProviderCoreTests/SlotSizingDriftTests.swift b/provider-swift/Tests/ProviderCoreTests/SlotSizingDriftTests.swift index 5cfa6b95b..0097fd17c 100644 --- a/provider-swift/Tests/ProviderCoreTests/SlotSizingDriftTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/SlotSizingDriftTests.swift @@ -94,8 +94,7 @@ private func gemma4Kinds(from config: [String: Any]) -> [CBv2LayerKind] { globalHeadDim: config["global_head_dim"] as! Int, numAttentionHeads: config["num_attention_heads"] as! Int, numKeyValueHeads: config["num_key_value_heads"] as! Int, - numGlobalKeyValueHeads: config["num_global_key_value_heads"] as? Int, - attentionKeqV: config["attention_k_eq_v"] as! Bool) + numGlobalKeyValueHeads: config["num_global_key_value_heads"] as? Int) } // MARK: - Tests @@ -191,8 +190,11 @@ struct SlotSizingDriftTests { // Engine truth from the REAL file. let engineRate: Int if checkpoint.textConfigWrapped { - let textConfig = try EngineV2VLMTextExtraction.decodeTextConfiguration( - configData: configData) + let root = try JSONSerialization.jsonObject(with: configData) as! [String: Any] + let textConfigData = try JSONSerialization.data( + withJSONObject: root["text_config"] as! [String: Any]) + let textConfig = try JSONDecoder().decode( + Gemma4TextConfiguration.self, from: textConfigData) engineRate = SlotSizingSnapshot.fp16KVBytesPerToken( layerKinds: textConfig.cbv2LayerKinds) } else { diff --git a/provider-swift/Tests/ProviderCoreTests/ThroughputSweepTests.swift b/provider-swift/Tests/ProviderCoreTests/ThroughputSweepTests.swift index 24f3e501d..634b4a786 100644 --- a/provider-swift/Tests/ProviderCoreTests/ThroughputSweepTests.swift +++ b/provider-swift/Tests/ProviderCoreTests/ThroughputSweepTests.swift @@ -103,6 +103,11 @@ struct DecodeBandwidthModelTests { @Suite("throughput sweep: row aggregation") struct ThroughputSweepRowAggregationTests { + @Test("decode sweep ignores EOS to preserve a fixed token budget") + func fixedDecodeBudget() { + #expect(ThroughputSweep.fixedBudgetStopTokens.isEmpty) + } + @Test("clean rows aggregate tokens and the slowest row's elapsed") func cleanRowsAggregate() { let cell = ThroughputSweep.aggregateRows([ @@ -207,7 +212,8 @@ struct ThroughputSweepReportTests { prefill: [.init(promptTokens: 128, prefillTokensPerSecond: 900, elapsedMs: 142)], decode: decodeSamples(b1Aggregate: 21), derived: derived, - notes: ["test"] + notes: ["test"], + gemmaOptimizations: .init(settings: GemmaOptimizationSettings()) ) let json = try report.jsonString() @@ -222,6 +228,13 @@ struct ThroughputSweepReportTests { #expect(decoded.modelID == report.modelID) #expect(decoded.decode.count == 2) #expect(decoded.derived.regime == .dense) + #expect(decoded.gemmaOptimizations.prefillLayer18) + #expect(decoded.gemmaOptimizations.weightedR1) + #expect(decoded.gemmaOptimizations.environment == [ + GemmaOptimizationEnvironment.prefillLayer18Key: "18", + GemmaOptimizationEnvironment.weightedUnsortKey: "1", + GemmaOptimizationEnvironment.safeR1Key: "1", + ]) #expect(decoded.schemaVersion == ThroughputSweepReport.currentSchemaVersion) } @@ -249,6 +262,7 @@ struct ThroughputSweepReportTests { decode: [], derived: derived, notes: ["kv backend: selection=paged, resolved=n/a (no decode cells ran)"], + gemmaOptimizations: .init(settings: GemmaOptimizationSettings()), decodeConstructionFailure: .init( kvBackendSelection: "paged", reason: reason) ) @@ -292,6 +306,7 @@ struct ThroughputSweepReportTests { decode: mixed, derived: derived, notes: [], + gemmaOptimizations: .init(settings: GemmaOptimizationSettings()), kvBackend: .init( selection: "auto", resolved: ["paged", "contiguous (fallback: pool capacity)"]) @@ -342,6 +357,7 @@ struct ThroughputSweepReportTests { ], derived: derived, notes: [], + gemmaOptimizations: .init(settings: GemmaOptimizationSettings()), kvBackend: .init(selection: "paged", resolved: ["paged"]), decodeCoverage: .init( requestedBatchSizes: [1, 8], @@ -379,6 +395,7 @@ struct ThroughputSweepReportTests { memoryBandwidthGbs: 546), prefill: [], decode: decodeSamples(b1Aggregate: 21), derived: derived, notes: [], + gemmaOptimizations: .init(settings: GemmaOptimizationSettings()), kvBackend: .init(selection: "paged", resolved: ["paged"]), decodeCoverage: .init(requestedBatchSizes: [1, 2], unmeasured: [])) let json = try report.jsonString() @@ -388,7 +405,7 @@ struct ThroughputSweepReportTests { ThroughputSweepReport.self, from: Data(json.utf8)) #expect(decoded.decodeCoverage.unmeasured.isEmpty) #expect(decoded.decodeCoverage.requestedBatchSizes == [1, 2]) - #expect(decoded.schemaVersion == 4) + #expect(decoded.schemaVersion == ThroughputSweepReport.currentSchemaVersion) } } diff --git a/scripts/fetch-metallib.sh b/scripts/fetch-metallib.sh index 70bb13786..5df8ed2d9 100755 --- a/scripts/fetch-metallib.sh +++ b/scripts/fetch-metallib.sh @@ -1,12 +1,10 @@ #!/bin/bash # fetch-metallib.sh -- build the matching mlx.metallib for local Swift builds. # -# NOTE: despite the name, this now BUILDS the .metallib from source rather than -# fetching it from a PyPI wheel. Building from our own fork guarantees the GPU -# kernels match the exact MLX commit the host C++ links against — including the -# resource-count trim and the M5 `_nax` kernels — and needs no published wheel -# (there is no mlx==0.32.0 on PyPI). This mirrors the release-swift.yml -# "Build mlx.metallib from source" step. +# NOTE: despite the name, this BUILDS the metallib from source rather than +# fetching a PyPI wheel. The kernels match the exact MLX source that the host +# C++ links against. This is the canonical metallib path for local, +# integration, CI, and release builds. # # mlx-swift's Cmlx target does NOT compile its Metal kernels through SwiftPM, so # we compile them here with cmake from libs/mlx-swift/Source/Cmlx/mlx (the same @@ -21,11 +19,20 @@ set -euo pipefail REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)" SWIFT_PROVIDER_DIR="${SWIFT_PROVIDER_DIR:-$REPO_ROOT/provider-swift}" -# Source of truth: the mlx submodule the Cmlx target actually compiles against. -MLX_SRC="${MLX_SRC:-$REPO_ROOT/libs/mlx-swift/Source/Cmlx/mlx}" +# Source of truth: exactly the mlx tree the Cmlx target compiles against. +MLX_SRC="$REPO_ROOT/libs/mlx-swift/Source/Cmlx/mlx" # The _nax kernels are only compiled when SDK >= 26.2 AND deployment target # >= 26.2 AND Metal >= 4.0 (mlx/backend/metal/kernels/CMakeLists.txt). DEPLOYMENT_TARGET="${MLX_METALLIB_DEPLOYMENT_TARGET:-26.2}" +JIT_MODE="OFF" +NAX_SYMBOL="_nax" +GEMV_SYMBOL="gemv" +R1_BUILDER_SYMBOL="build_gemma4_sorted_expert_tiles_bm32" +R1_KERNEL_SYMBOL="affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_alN_true_bm_32_bn_32_bk_32" +COMPLETENESS_CONTRACT="$( + printf '%s\n' "$NAX_SYMBOL" "$GEMV_SYMBOL" "$R1_BUILDER_SYMBOL" "$R1_KERNEL_SYMBOL" \ + | shasum -a 256 | cut -d' ' -f1 +)" TARGET_ARG="${1:-debug}" case "$TARGET_ARG" in @@ -37,6 +44,8 @@ esac mkdir -p "$DEST_DIR" command -v cmake >/dev/null 2>&1 || { echo "✗ cmake not found (brew install cmake)"; exit 1; } +command -v xcodebuild >/dev/null 2>&1 || { echo "✗ xcodebuild not found"; exit 1; } +command -v xcrun >/dev/null 2>&1 || { echo "✗ xcrun not found"; exit 1; } test -f "$MLX_SRC/mlx/version.h" || { echo "✗ mlx submodule missing at $MLX_SRC" echo " run: git submodule update --init --recursive" @@ -54,38 +63,51 @@ test -f "$MLX_SRC/mlx/version.h" || { # change the hash instead of collapsing to "Binary files differ", # --no-ext-diff so a user's difftool cannot alter the key, and --no-color so # terminal settings cannot either. -if git -C "$MLX_SRC" rev-parse --is-inside-work-tree >/dev/null 2>&1; then - MLX_SHA="$(git -C "$MLX_SRC" rev-parse HEAD 2>/dev/null || echo nogit)" - MLX_TREE_HASH="$({ - git -C "$MLX_SRC" diff HEAD --binary --no-ext-diff --no-color -- 2>/dev/null || true - { - git -C "$MLX_SRC" ls-files --others --exclude-standard 2>/dev/null || true - } | while IFS= read -r f; do - printf '%s\n' "$f" - cat "$MLX_SRC/$f" 2>/dev/null || true - done - } | shasum -a 256 | cut -d' ' -f1)" -else - # Not a git checkout (source tarball, vendored copy). There is no commit - # and no diff to read, so hashing nothing would mark ANY local edit as - # "clean" and reintroduce exactly the bug this key is meant to prevent. - # Hash the source contents instead. - MLX_SHA="nogit" - MLX_TREE_HASH="$( - find "$MLX_SRC/mlx" -type f \( -name '*.h' -o -name '*.metal' -o -name '*.cpp' \) -print0 2>/dev/null \ - | sort -z | xargs -0 shasum -a 256 2>/dev/null | shasum -a 256 | cut -d' ' -f1 - )" +compute_mlx_identity() { + if git -C "$MLX_SRC" rev-parse --is-inside-work-tree >/dev/null 2>&1; then + MLX_SHA="$(git -C "$MLX_SRC" rev-parse HEAD 2>/dev/null || echo nogit)" + MLX_TREE_HASH="$({ + git -C "$MLX_SRC" diff HEAD --binary --no-ext-diff --no-color -- 2>/dev/null || true + { + git -C "$MLX_SRC" ls-files --others --exclude-standard 2>/dev/null || true + } | while IFS= read -r f; do + printf '%s\n' "$f" + cat "$MLX_SRC/$f" 2>/dev/null || true + done + } | shasum -a 256 | cut -d' ' -f1)" + else + # Hash the entire source tree (including CMake inputs and generated + # manifests), not just kernel extensions. + MLX_SHA="nogit" + MLX_TREE_HASH="$( + find "$MLX_SRC" -type f ! -path '*/.git/*' -print0 2>/dev/null \ + | sort -z | xargs -0 shasum -a 256 2>/dev/null | shasum -a 256 | cut -d' ' -f1 + )" + fi +} + +compute_mlx_identity +INITIAL_MLX_SHA="$MLX_SHA" +INITIAL_MLX_TREE_HASH="$MLX_TREE_HASH" +if [ "$MLX_SHA" = "nogit" ]; then echo "→ mlx source is not a git checkout; keying metallib on file contents (${MLX_TREE_HASH:0:12})" fi +assert_mlx_source_unchanged() { + compute_mlx_identity + if [ "$MLX_SHA" != "$INITIAL_MLX_SHA" ] \ + || [ "$MLX_TREE_HASH" != "$INITIAL_MLX_TREE_HASH" ] + then + echo "✗ MLX source changed while preparing metallib; refusing stale cache publication" + return 1 + fi +} # Hash of empty input == a clean git tree. MLX_CLEAN_HASH="$(printf '' | shasum -a 256 | cut -d' ' -f1)" if [ "$MLX_SHA" != "nogit" ] && [ "$MLX_TREE_HASH" = "$MLX_CLEAN_HASH" ]; then - # Clean tree. The key is versioned (-c2) rather than bare: entries written - # by the PREVIOUS version of this script used the bare name and could have - # been built from a DIRTY tree, so reusing them would resurrect the very - # mismatch this change exists to prevent. The suffix retires them. - MLX_TREE_SUFFIX="-c2" + # Keep a visible clean-tree epoch; the helper and completeness contract are + # independently folded into the toolchain hash below. + MLX_TREE_SUFFIX="-c4" else MLX_TREE_SUFFIX="-w${MLX_TREE_HASH:0:12}" if [ "$MLX_SHA" != "nogit" ]; then @@ -93,40 +115,178 @@ else fi fi -# Cache the built metallib by mlx commit + working-tree contents + deployment -# target so repeat runs are instant (the build is ~1 min). Override with -# METALLIB_CACHE_DIR. +# Cache identity covers every input which can change generated Metal code or +# the helper's acceptance contract. In particular, a cache from another Xcode +# or SDK must never be reused merely because the source commit is identical. +XCODE_VERSION="$(xcodebuild -version | tr '\n' ';')" +SDK_VERSION="$(xcrun --sdk macosx --show-sdk-version)" +SDK_BUILD_VERSION="$(xcrun --sdk macosx --show-sdk-build-version)" +HELPER_CONTRACT_HASH="$(shasum -a 256 "$0" | cut -d' ' -f1)" +TOOLCHAIN_HASH="$( + printf '%s\n' \ + "$XCODE_VERSION" \ + "$SDK_VERSION" \ + "$SDK_BUILD_VERSION" \ + "deployment=$DEPLOYMENT_TARGET" \ + "jit=$JIT_MODE" \ + "helper=$HELPER_CONTRACT_HASH" \ + "completeness=$COMPLETENESS_CONTRACT" \ + | shasum -a 256 | cut -d' ' -f1 +)" + CACHE_DIR="${METALLIB_CACHE_DIR:-/tmp/mlx-metallib-cache}" -CACHE_KEY="mlx-${MLX_SHA}${MLX_TREE_SUFFIX}-dt${DEPLOYMENT_TARGET}" +CACHE_KEY="mlx-${MLX_SHA}${MLX_TREE_SUFFIX}-tc${TOOLCHAIN_HASH:0:16}-dt${DEPLOYMENT_TARGET}-jitoff-c${COMPLETENESS_CONTRACT:0:12}" CACHED="$CACHE_DIR/${CACHE_KEY}.metallib" +verify_metallib() { + local metallib="$1" + local symbol matches + + if [ ! -s "$metallib" ]; then + echo "✗ metallib missing or empty: $metallib" + return 1 + fi + + for symbol in \ + "$NAX_SYMBOL" \ + "$GEMV_SYMBOL" \ + "$R1_BUILDER_SYMBOL" \ + "$R1_KERNEL_SYMBOL" + do + # Use grep -c rather than grep -q: grep -q closes the pipe after its + # first match, causing strings to receive SIGPIPE under pipefail. + matches="$(strings "$metallib" | grep -F -c "$symbol" || true)" + if [ "$matches" -eq 0 ]; then + echo "✗ metallib missing required symbol string: $symbol" + return 1 + fi + done +} + +mkdir -p "$CACHE_DIR" + +# Serialize validation and publication for a cache key. A symlink is the +# atomic owner-bearing primitive: it never exposes a lock without PID/start +# identity. The unique target directory also provides an atomic reaper claim. +CACHE_LOCK="$CACHE_DIR/.${CACHE_KEY}.lock" +PROCESS_START_HASH="$( + LC_ALL=C TZ=UTC ps -p $$ -o lstart= | shasum -a 256 | cut -d' ' -f1 +)" +LOCK_STATE_NAME=".${CACHE_KEY}.owner-$$-$PROCESS_START_HASH" +LOCK_STATE="$CACHE_DIR/$LOCK_STATE_NAME" +mkdir "$LOCK_STATE" +LOCK_HELD=0 +LOCK_ATTEMPTS=0 + +while :; do + if ln -s "$LOCK_STATE_NAME" "$CACHE_LOCK" 2>/dev/null; then + LOCK_HELD=1 + break + fi + + lock_target="$(readlink "$CACHE_LOCK" 2>/dev/null || true)" + case "$lock_target" in + ".${CACHE_KEY}.owner-"*) + owner_meta="${lock_target#".${CACHE_KEY}.owner-"}" + owner_pid="${owner_meta%%-*}" + owner_start_hash="${owner_meta#*-}" + current_start_hash="" + if kill -0 "$owner_pid" 2>/dev/null; then + current_start_hash="$( + LC_ALL=C TZ=UTC ps -p "$owner_pid" -o lstart= 2>/dev/null \ + | shasum -a 256 | cut -d' ' -f1 + )" + fi + if [ "$current_start_hash" != "$owner_start_hash" ]; then + stale_state="$CACHE_DIR/$lock_target" + # Recreate a missing target left by an interrupted release. + if [ ! -d "$stale_state" ]; then + mkdir "$stale_state" 2>/dev/null || true + fi + if mkdir "$stale_state/reaper" 2>/dev/null; then + if [ "$(readlink "$CACHE_LOCK" 2>/dev/null || true)" = "$lock_target" ]; then + rm -f "$CACHE_LOCK" + fi + rmdir "$stale_state/reaper" + rmdir "$stale_state" + rm -rf "$CACHE_DIR"/.build-"$CACHE_KEY"."$owner_pid".* + rm -f "$CACHE_DIR/.${CACHE_KEY}.${owner_pid}" + continue + fi + fi + ;; + esac + + if [ "$LOCK_ATTEMPTS" -eq 0 ]; then + echo "→ Waiting for concurrent metallib build $CACHE_KEY" + fi + LOCK_ATTEMPTS=$((LOCK_ATTEMPTS + 1)) + if [ "$LOCK_ATTEMPTS" -ge 6000 ]; then + echo "✗ timed out waiting for metallib cache lock: $CACHE_LOCK" + exit 1 + fi + sleep 0.1 +done + +BUILD_DIR="" +CACHE_TMP="" +DEST_TMP="" +release_cache_lock() { + if [ "$LOCK_HELD" -eq 1 ]; then + if [ "$(readlink "$CACHE_LOCK" 2>/dev/null || true)" = "$LOCK_STATE_NAME" ]; then + rm -f "$CACHE_LOCK" + fi + LOCK_HELD=0 + fi + rmdir "$LOCK_STATE" 2>/dev/null || true +} +cleanup() { + [ -z "$BUILD_DIR" ] || rm -rf "$BUILD_DIR" + [ -z "$CACHE_TMP" ] || rm -f "$CACHE_TMP" + [ -z "$DEST_TMP" ] || rm -f "$DEST_TMP" + release_cache_lock +} +trap cleanup EXIT +trap 'exit 143' HUP INT TERM + +if [ -s "$CACHED" ]; then + if verify_metallib "$CACHED"; then + echo "→ Using cached metallib $CACHE_KEY" + else + echo "→ Discarding incomplete cached metallib $CACHE_KEY" + rm -f "$CACHED" + fi +fi + if [ ! -s "$CACHED" ]; then echo "→ Building mlx.metallib from $MLX_SRC @ $CACHE_KEY" - BUILD_DIR="$(mktemp -d)" + BUILD_DIR="$(mktemp -d "$CACHE_DIR/.build-${CACHE_KEY}.$$.XXXXXX")" + CACHE_TMP="$CACHE_DIR/.${CACHE_KEY}.$$" + cmake -S "$MLX_SRC" -B "$BUILD_DIR" \ -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_DEPLOYMENT_TARGET="$DEPLOYMENT_TARGET" \ - -DMLX_METAL_JIT=OFF \ + -DMLX_METAL_JIT="$JIT_MODE" \ -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_EXAMPLES=OFF \ -DMLX_BUILD_BENCHMARKS=OFF -DMLX_BUILD_PYTHON_BINDINGS=OFF >/dev/null cmake --build "$BUILD_DIR" --target mlx-metallib -j"$(sysctl -n hw.ncpu 2>/dev/null || echo 4)" - mkdir -p "$CACHE_DIR" - cp "$BUILD_DIR/mlx/backend/metal/kernels/mlx.metallib" "$CACHED" + cp "$BUILD_DIR/mlx/backend/metal/kernels/mlx.metallib" "$CACHE_TMP" + verify_metallib "$CACHE_TMP" + assert_mlx_source_unchanged + mv -f "$CACHE_TMP" "$CACHED" + CACHE_TMP="" rm -rf "$BUILD_DIR" -else - echo "→ Using cached metallib $CACHE_KEY" -fi - -# Sanity: the _nax kernels must be present, otherwise the build silently used an -# SDK/deployment target < 26.2 and we'd ship a nax-less metallib. Use `grep -c` -# (not `grep -q`): grep -q closes the pipe on first match, which makes `strings` -# on this 150 MB+ file exit via SIGPIPE and trips `set -o pipefail` — a false -# negative. -NAX_KERNELS="$(strings "$CACHED" | grep -c "_nax" || true)" -if [ "$NAX_KERNELS" -eq 0 ]; then - echo "✗ built metallib has no _nax kernels — is your Xcode SDK / deployment target >= 26.2?" - exit 1 + BUILD_DIR="" fi -cp "$CACHED" "$DEST_DIR/mlx.metallib" +# Always replace the destination, even on a cache hit. A merely existing file +# may have been built from different host sources and can hang the GPU. +DEST_TMP="$DEST_DIR/.mlx.metallib.$$" +assert_mlx_source_unchanged +cp "$CACHED" "$DEST_TMP" +mv -f "$DEST_TMP" "$DEST_DIR/mlx.metallib" +DEST_TMP="" echo "✓ wrote $DEST_DIR/mlx.metallib ($(shasum -a 256 "$DEST_DIR/mlx.metallib" | cut -d' ' -f1))" + +release_cache_lock +trap - EXIT HUP INT TERM diff --git a/scripts/gemma_contbatch/backend.py b/scripts/gemma_contbatch/backend.py index 1c6513e56..46f39225d 100644 --- a/scripts/gemma_contbatch/backend.py +++ b/scripts/gemma_contbatch/backend.py @@ -2,10 +2,11 @@ A decode curve is only comparable to another decode curve if both were produced by the same KV backend. Nothing else in this harness can establish -that: `--kv-backend auto` is a *selection*. It resolves paged as of v0.8.0 -(see the provider's `EngineV2Factory.prepareProductionBackend`) but degrades -to contiguous on a box that cannot serve paged, and an explicit `paged` can -still be vetoed by the fleet kill switch. A run that +that: `--kv-backend auto` is a *selection* and resolves contiguous as of +v0.8.1 (see the provider's `EngineV2Factory.prepareProductionBackend`). The +wrapper requests `paged` explicitly; that selection refuses when paged cannot +be built, while the fleet kill switch can deliberately degrade it to +contiguous. A run that did not build the backend it names measures the fallback while every other check in this wrapper stays green, and a percentage delta against a baseline recorded on the other backend is a backend change wearing a performance diff --git a/scripts/gemma_contbatch/baseline.py b/scripts/gemma_contbatch/baseline.py index eb6eb8f73..7aa6c52a7 100644 --- a/scripts/gemma_contbatch/baseline.py +++ b/scripts/gemma_contbatch/baseline.py @@ -15,6 +15,7 @@ from .config import SCHEMA_VERSION from .environment import baseline_environment +from .gemma_optimizations import validate_gemma_optimizations NO_COMPARE_HINT = "omit --baseline to run without a comparison" @@ -101,11 +102,18 @@ def validate_hardware_pin(baseline: dict, hardware: dict) -> None: ) -def validate_configuration_pin(args: argparse.Namespace, baseline: dict) -> None: +def validate_configuration_pin( + args: argparse.Namespace, + baseline: dict, + gemma_optimizations: dict, + comparison_axis: str = "code", +) -> None: baseline_configuration = baseline.get("configuration") if not isinstance(baseline_configuration, dict): raise RuntimeError("baseline does not record a configuration; " + NO_COMPARE_HINT) expected = { + "iterations": args.iterations, + "decodeIterations": args.iterations, "decodePromptTokens": args.decode_prompt_tokens, "decodeTokens": args.decode_tokens, "arrivalPromptTokens": args.arrival_prompt_tokens, @@ -125,6 +133,41 @@ def validate_configuration_pin(args: argparse.Namespace, baseline: dict) -> None + "; " + NO_COMPARE_HINT ) + baseline_gemma = validate_gemma_optimizations( + baseline_configuration.get("gemmaOptimizations"), "baseline configuration" + ) + if comparison_axis == "code": + if baseline_gemma != gemma_optimizations: + raise RuntimeError( + "gemmaOptimizations differ on a code comparison; " + NO_COMPARE_HINT + ) + elif comparison_axis == "gemma-optimizations": + if baseline_gemma == gemma_optimizations: + raise RuntimeError( + "gemma-optimizations comparison requires different effective settings; " + + NO_COMPARE_HINT + ) + else: + raise RuntimeError(f"unknown comparison axis {comparison_axis!r}") + + +def validate_artifact_pin(baseline: dict, metadata: dict) -> None: + baseline_metadata = baseline.get("metadata") + if not isinstance(baseline_metadata, dict): + raise RuntimeError("baseline does not record artifact hashes; " + NO_COMPARE_HINT) + mismatches = [ + f"{key}={metadata.get(key)!r} (baseline {baseline_metadata.get(key)!r})" + for key in ("binarySha256", "metallibSha256") + if metadata.get(key) != baseline_metadata.get(key) + ] + if mismatches: + raise RuntimeError( + "gemma-optimizations comparison requires identical artifacts: " + + ", ".join(mismatches) + + "; " + + NO_COMPARE_HINT + ) + def describe_env_value(value: str | None) -> str: """`unset` is a bareword so it cannot be confused with the string 'unset'.""" @@ -188,11 +231,9 @@ def validate_environment_pin(baseline: dict, environment: dict[str, str]) -> Non def validate_schema_version_pin(baseline: dict) -> None: """Refuse a baseline written against a different wrapper schema. - Schema 3 removed `configuration.maxBatch` and added the `kvBackend` - block. A schema-2 baseline therefore records a batch ladder this runner - cannot see and no backend at all, so every pin below it reads absent - fields as "not recorded" and the comparison silently comes out as a - same-shape delta between two different experiments. Fail here instead. + Schema 5 additionally records the config-projected Gemma posture. Older + reports cannot distinguish an ON run from an OFF run, so every named field + below must come from the exact schema this runner writes. """ recorded = baseline.get("schemaVersion") if recorded == SCHEMA_VERSION: @@ -200,8 +241,8 @@ def validate_schema_version_pin(baseline: dict) -> None: raise RuntimeError( f"baseline schemaVersion is {recorded!r}, this runner writes " f"{SCHEMA_VERSION}; the two reports do not describe the same fields " - f"(schema 3 replaced configuration.maxBatch with configuration." - f"batchSizes and added the kvBackend block); " + f"(schema 5 includes batchSizes, kvBackend, and effective Gemma " + f"optimization settings); " + "re-record the baseline with this runner, or " + NO_COMPARE_HINT ) @@ -212,6 +253,9 @@ def validate_baseline_pins( model_snapshot: str, hardware: dict, environment: dict[str, str], + gemma_optimizations: dict, + comparison_axis: str = "code", + metadata: dict | None = None, ) -> None: """Every pin that must hold before a delta can be read as an engine delta.""" # First: every pin below reads named fields, and a schema mismatch makes @@ -219,5 +263,9 @@ def validate_baseline_pins( validate_schema_version_pin(baseline) validate_model_pin(baseline, args.model, model_snapshot) validate_hardware_pin(baseline, hardware) - validate_configuration_pin(args, baseline) + validate_configuration_pin( + args, baseline, gemma_optimizations, comparison_axis=comparison_axis + ) + if comparison_axis == "gemma-optimizations": + validate_artifact_pin(baseline, metadata or {}) validate_environment_pin(baseline, environment) diff --git a/scripts/gemma_contbatch/config.py b/scripts/gemma_contbatch/config.py index 2a85a8341..92d148925 100644 --- a/scripts/gemma_contbatch/config.py +++ b/scripts/gemma_contbatch/config.py @@ -20,27 +20,24 @@ # not just the decode curve's. The scheduler-prefill and arrival # commands now take the selection too, so `kvBackend.resolved` is the # whole run's population rather than the sweep's. -SCHEMA_VERSION = 4 +# 5 — required effective config-projected Gemma settings, validated across +# all three subprocesses and pinned for baseline comparisons. +SCHEMA_VERSION = 5 DEFAULT_MODEL = "mlx-community/gemma-4-26B-A4B-it-qat-4bit" # The canonical posture this release is measured under. Both defaults are # load-bearing: # -# paged — `auto` resolves PAGED as of v0.8.0, but it degrades -# SILENTLY (kill switch, kernel preflight, pool capacity) -# while an explicit `paged` REFUSES. Naming the backend is -# the only way this wrapper can promise it measured what it -# reports, so it is requested by name even though `auto` -# would usually land on the same engine. -# 1,2,4,8 — paged-vs-contiguous aggregate decode crosses over at ~B=5 -# (measured on gemma-4 / M4 Max: 0.92x at B=1, 0.98x at B=4, -# 1.17x at B=8). A curve that stops at 4 structurally cannot -# observe the win and reads as a regression. B=8 is also the -# raised production concurrency ceiling. The list is sparse on +# contiguous — `auto` resolves contiguous as of v0.8.1, but naming the +# backend keeps every phase's release posture explicit and +# prevents a future default flip from changing the benchmark. +# 1,2,4,8 — the current production concurrency default is B=4 under the +# contiguous `auto` posture, while B=8 remains a supported +# stress point. The list is sparse on # purpose: a dense 1..8 ladder doubles wall time for cells no # gate reads. -DEFAULT_KV_BACKEND = "paged" +DEFAULT_KV_BACKEND = "contiguous" DEFAULT_BATCH_SIZES = [1, 2, 4, 8] KV_BACKENDS = ("auto", "contiguous", "paged") EXPECTED_ARRIVAL_PATTERNS = { @@ -69,6 +66,21 @@ def parse_args() -> argparse.Namespace: ) ) parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument( + "--config", + default=None, + help="provider TOML passed to every darkbloom benchmark subprocess", + ) + parser.add_argument( + "--comparison-axis", + choices=("code", "gemma-optimizations"), + default="code", + help=( + "dimension allowed to differ from --baseline: code requires equal " + "Gemma settings; gemma-optimizations requires equal binaries and " + "different effective Gemma settings" + ), + ) parser.add_argument("--iterations", type=int, default=3) parser.add_argument( "--prefill-lengths", type=parse_positive_ints, default=[128, 512, 2048] @@ -87,7 +99,7 @@ def parse_args() -> argparse.Namespace: choices=KV_BACKENDS, default=DEFAULT_KV_BACKEND, help=( - "KV backend the decode sweep is built with " + "KV backend every benchmark phase is built with " f"(default {DEFAULT_KV_BACKEND}; 'auto' may silently resolve " "either backend, so it cannot pin a measurement)" ), diff --git a/scripts/gemma_contbatch/environment.py b/scripts/gemma_contbatch/environment.py index efc804544..1f0871426 100644 --- a/scripts/gemma_contbatch/environment.py +++ b/scripts/gemma_contbatch/environment.py @@ -58,6 +58,7 @@ "MLX_COMPILED_DECODE", "MLX_DISABLE_COMPILE", "MLX_GATHER_QMM_EXPERT_SLICES", + "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT", "MLX_METALLIB_PATH", } ) diff --git a/scripts/gemma_contbatch/gemma_optimizations.py b/scripts/gemma_contbatch/gemma_optimizations.py new file mode 100644 index 000000000..0eafafdf4 --- /dev/null +++ b/scripts/gemma_contbatch/gemma_optimizations.py @@ -0,0 +1,50 @@ +"""Effective Gemma optimization provenance across benchmark subprocesses.""" + +from __future__ import annotations + +from collections.abc import Mapping + + +PREFILL_KEY = "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL" +WEIGHTED_KEY = "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT" +SAFE_R1_KEY = "MLX_GATHER_QMM_EXPERT_SLICES" + + +def validate_gemma_optimizations(value: object, source: str) -> dict: + if not isinstance(value, Mapping): + raise RuntimeError(f"{source} did not report effective Gemma optimizations") + prefill = value.get("prefillLayer18") + weighted = value.get("weightedR1") + environment = value.get("environment") + if not isinstance(prefill, bool) or not isinstance(weighted, bool): + raise RuntimeError(f"{source} reported malformed Gemma optimization booleans") + expected_environment = { + PREFILL_KEY: "18" if prefill else "0", + WEIGHTED_KEY: "1" if weighted else "0", + SAFE_R1_KEY: "1" if weighted else "0", + } + if environment != expected_environment: + raise RuntimeError( + f"{source} Gemma environment projection does not match its settings: " + f"expected {expected_environment}, got {environment!r}" + ) + return { + "prefillLayer18": prefill, + "weightedR1": weighted, + "environment": expected_environment, + } + + +def resolve_gemma_optimizations(raw_outputs: Mapping[str, dict]) -> dict: + """Require one valid effective posture shared by every benchmark phase.""" + resolved = { + name: validate_gemma_optimizations(payload.get("gemmaOptimizations"), name) + for name, payload in raw_outputs.items() + } + distinct = {repr(value) for value in resolved.values()} + if len(distinct) != 1: + detail = ", ".join( + f"{name}={value}" for name, value in sorted(resolved.items()) + ) + raise RuntimeError(f"benchmark phases used different Gemma optimizations: {detail}") + return next(iter(resolved.values())) diff --git a/scripts/gemma_contbatch/report.py b/scripts/gemma_contbatch/report.py index 1e9264c68..a89be70cc 100644 --- a/scripts/gemma_contbatch/report.py +++ b/scripts/gemma_contbatch/report.py @@ -163,6 +163,10 @@ def markdown_report(report: dict) -> str: f"| Decode batch sizes | {', '.join(map(str, configuration['batchSizes']))} |", f"| Decode samples per batch | {configuration['decodeIterations']} |", f"| Arrival prompt / output | {configuration['arrivalPromptTokens']} / {configuration['arrivalDecodeTokens']} tokens |", + f"| Provider config | `{configuration['providerConfig']}` |", + f"| Comparison axis | `{configuration['comparisonAxis']}` |", + f"| Gemma layer-18 prefill | {'on' if configuration['gemmaOptimizations']['prefillLayer18'] else 'off'} |", + f"| Gemma weighted-unsort + safe-R1 | {'on' if configuration['gemmaOptimizations']['weightedR1'] else 'off'} |", ] lines += kv_backend_lines(kv_backend) diff --git a/scripts/gemma_contbatch/runner.py b/scripts/gemma_contbatch/runner.py index c8d440f58..3a8991f96 100644 --- a/scripts/gemma_contbatch/runner.py +++ b/scripts/gemma_contbatch/runner.py @@ -20,6 +20,7 @@ from .checks import assert_finite from .config import SCHEMA_VERSION, parse_args from .environment import performance_environment +from .gemma_optimizations import resolve_gemma_optimizations from .process import ( BenchmarkCommandFailure, atomic_write, @@ -49,6 +50,30 @@ def resolve_output_dir(args: argparse.Namespace, repo_root: Path) -> Path: return output_dir +def resolve_config_path(raw_path: str | None, repo_root: Path) -> str | None: + """Normalize an explicit config against repo root and fail if unreadable.""" + if raw_path is None: + return None + path = Path(raw_path).expanduser() + if not path.is_absolute(): + path = repo_root / path + try: + path = path.resolve(strict=True) + except OSError as error: + raise RuntimeError(f"provider config is not readable: {path}") from error + if not path.is_file() or not os.access(path, os.R_OK): + raise RuntimeError(f"provider config is not a readable regular file: {path}") + return str(path) + + +def benchmark_argv(binary: Path, args: argparse.Namespace) -> list[str]: + """Shared prefix so every phase loads the same provider config.""" + command = [str(binary), "benchmark"] + if args.config: + command.extend(["--config", args.config]) + return command + ["--model", args.model] + + def persist_failed_report(failure: BenchmarkCommandFailure, output_dir: Path) -> int: """Keep the structured report a failed benchmark deliberately printed. @@ -157,6 +182,7 @@ def arrival_argv(benchmark: list[str], args: argparse.Namespace) -> list[str]: def main() -> int: args = parse_args() repo_root = Path(__file__).resolve().parents[2] + args.config = resolve_config_path(args.config, repo_root) provider_dir = repo_root / "provider-swift" binary = provider_dir / ".build/release/darkbloom" metallib = provider_dir / ".build/release/mlx.metallib" @@ -193,7 +219,7 @@ def main() -> int: raise RuntimeError("release binary or mlx.metallib is missing") output_dir = resolve_output_dir(args, repo_root) - benchmark = [str(binary), "benchmark", "--model", args.model] + benchmark = benchmark_argv(binary, args) # Parse before aborting: a refused sweep prints its report and THEN # fails, and that report is the whole diagnostic. try: @@ -209,6 +235,7 @@ def main() -> int: "arrivalInvariance": arrival, } validate_raw_outputs(args, sweep, scheduler, arrival) + gemma_optimizations = resolve_gemma_optimizations(raw_outputs) # The backend every phase was actually built with. Extracted before the # summary so a run that cannot name its backends never reaches a report, # let alone a comparison. @@ -220,6 +247,18 @@ def main() -> int: # The same capture is recorded in the report and pinned against the # baseline below, so the two can never describe different runs. environment = performance_environment(os.environ) + metadata = { + "rootCommit": capture(["git", "rev-parse", "HEAD"], repo_root), + "gitStatus": capture(["git", "status", "--short"], repo_root).splitlines(), + "submodules": capture(["git", "submodule", "status"], repo_root).splitlines(), + "mlxMetalSourceStatus": mlx_metal_status, + "mlxMetalSourceFingerprint": mlx_metal_fingerprint, + "binarySha256": sha256(binary), + "metallibSha256": sha256(metallib), + "environment": environment, + "thermalBefore": thermal_before, + "thermalAfter": capture_optional(["pmset", "-g", "therm"], repo_root), + } report = { # See config.SCHEMA_VERSION for what each version changed. "schemaVersion": SCHEMA_VERSION, @@ -243,19 +282,11 @@ def main() -> int: "decodeIterations": args.iterations, "arrivalPromptTokens": args.arrival_prompt_tokens, "arrivalDecodeTokens": args.arrival_decode_tokens, + "providerConfig": args.config or "default", + "comparisonAxis": args.comparison_axis, + "gemmaOptimizations": gemma_optimizations, }, - "metadata": { - "rootCommit": capture(["git", "rev-parse", "HEAD"], repo_root), - "gitStatus": capture(["git", "status", "--short"], repo_root).splitlines(), - "submodules": capture(["git", "submodule", "status"], repo_root).splitlines(), - "mlxMetalSourceStatus": mlx_metal_status, - "mlxMetalSourceFingerprint": mlx_metal_fingerprint, - "binarySha256": sha256(binary), - "metallibSha256": sha256(metallib), - "environment": environment, - "thermalBefore": thermal_before, - "thermalAfter": capture_optional(["pmset", "-g", "therm"], repo_root), - }, + "metadata": metadata, "summary": summary, "raw": raw_outputs, "durationSeconds": time.monotonic() - started, @@ -271,10 +302,19 @@ def main() -> int: # delta is computed: an unpinned snapshot, a different Mac, or a # flipped kill switch turns unrelated differences into what reads as # an engine regression. - validate_baseline_pins(args, baseline, model_snapshot, hardware, environment) - # The pin this release turns on. A paged-versus-contiguous difference - # would otherwise show up as a double-digit aggregate delta with no - # trace of its cause anywhere in the report. + validate_baseline_pins( + args, + baseline, + model_snapshot, + hardware, + environment, + gemma_optimizations, + comparison_axis=args.comparison_axis, + metadata=metadata, + ) + # Keep the release's contiguous posture fixed across comparisons. A + # paged-versus-contiguous difference would otherwise read as a code or + # Gemma-optimization delta with no trace of its cause in the summary. validate_kv_backend_pin(baseline, kv_backend) report["comparison"] = compare(summary, baseline) diff --git a/scripts/gemma_contbatch/tests/fixtures.py b/scripts/gemma_contbatch/tests/fixtures.py index f5394e237..1e155cda2 100644 --- a/scripts/gemma_contbatch/tests/fixtures.py +++ b/scripts/gemma_contbatch/tests/fixtures.py @@ -8,6 +8,8 @@ from __future__ import annotations +import copy + from statistics import median from ..config import EXPECTED_ARRIVAL_PATTERNS @@ -21,6 +23,15 @@ } MODEL_ID = "mlx-community/gemma-4-26B-A4B-it-qat-4bit" MODEL_PATH = "/models/hub/models--mlx-community--gemma/snapshots/abc123" +GEMMA_OPTIMIZATIONS = { + "prefillLayer18": True, + "weightedR1": True, + "environment": { + "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL": "18", + "MLX_GEMMA4_FUSED_WEIGHTED_UNSORT": "1", + "MLX_GATHER_QMM_EXPERT_SLICES": "1", + }, +} def sweep_payload( @@ -71,10 +82,11 @@ def sweep_payload( } ) return { - "schemaVersion": 4, + "schemaVersion": 5, "modelID": MODEL_ID, "modelPath": MODEL_PATH, "hardware": dict(HARDWARE), + "gemmaOptimizations": copy.deepcopy(GEMMA_OPTIMIZATIONS), "prefill": prefill, "decode": decode, "derived": { @@ -105,9 +117,10 @@ def scheduler_payload( the binary emits it. """ return { - "schemaVersion": 1, + "schemaVersion": 2, "modelID": MODEL_ID, "modelPath": MODEL_PATH, + "gemmaOptimizations": copy.deepcopy(GEMMA_OPTIMIZATIONS), "kvBackend": {"selection": selection, "resolved": [resolved]}, "samples": [ { @@ -186,7 +199,8 @@ def arrival_payload( return { "modelID": MODEL_ID, "modelPath": MODEL_PATH, - "schemaVersion": 3, + "schemaVersion": 4, + "gemmaOptimizations": copy.deepcopy(GEMMA_OPTIMIZATIONS), "kvBackend": {"selection": selection, "resolved": [resolved]}, "promptTokensPerRequest": prompt_tokens, "decodeTokensPerRequest": decode_tokens, diff --git a/scripts/gemma_contbatch/tests/test_gemma_optimizations.py b/scripts/gemma_contbatch/tests/test_gemma_optimizations.py new file mode 100644 index 000000000..1c2f98613 --- /dev/null +++ b/scripts/gemma_contbatch/tests/test_gemma_optimizations.py @@ -0,0 +1,214 @@ +"""Gemma config provenance must survive every subprocess boundary.""" + +from __future__ import annotations + +import copy +import tempfile +import unittest +from pathlib import Path + +from .. import runner +from ..baseline import validate_artifact_pin, validate_configuration_pin +from ..gemma_optimizations import resolve_gemma_optimizations +from ..validation import validate_raw_outputs +from . import fixtures +from .test_kv_backend import make_args, make_arrival, make_scheduler, make_sweep + + +def raw_outputs() -> dict[str, dict]: + return { + "throughputSweep": make_sweep(), + "schedulerPrefill": make_scheduler(), + "arrivalInvariance": make_arrival(), + } + + +def baseline_configuration(args) -> dict: + return { + "iterations": args.iterations, + "decodeIterations": args.iterations, + "decodePromptTokens": args.decode_prompt_tokens, + "decodeTokens": args.decode_tokens, + "arrivalPromptTokens": args.arrival_prompt_tokens, + "arrivalDecodeTokens": args.arrival_decode_tokens, + "batchSizes": args.batch_sizes, + "prefillLengths": args.prefill_lengths, + "gemmaOptimizations": copy.deepcopy(fixtures.GEMMA_OPTIMIZATIONS), + } + + +def opposite_gemma_optimizations() -> dict: + settings = copy.deepcopy(fixtures.GEMMA_OPTIMIZATIONS) + settings["prefillLayer18"] = False + settings["environment"]["DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL"] = "0" + return settings + + +class GemmaOptimizationProvenanceTests(unittest.TestCase): + def test_matching_effective_settings_are_resolved(self): + self.assertEqual( + resolve_gemma_optimizations(raw_outputs()), fixtures.GEMMA_OPTIMIZATIONS + ) + + def test_phase_mismatch_is_refused(self): + outputs = raw_outputs() + outputs["arrivalInvariance"]["gemmaOptimizations"] = { + "prefillLayer18": False, + "weightedR1": True, + "environment": { + **fixtures.GEMMA_OPTIMIZATIONS["environment"], + "DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL": "0", + }, + } + with self.assertRaisesRegex(RuntimeError, "different Gemma optimizations"): + resolve_gemma_optimizations(outputs) + + def test_partial_weighted_projection_is_refused(self): + outputs = raw_outputs() + outputs["throughputSweep"]["gemmaOptimizations"]["environment"][ + "MLX_GATHER_QMM_EXPERT_SLICES" + ] = "0" + with self.assertRaisesRegex(RuntimeError, "projection does not match"): + resolve_gemma_optimizations(outputs) + + def test_stale_raw_schema_is_refused(self): + outputs = raw_outputs() + outputs["throughputSweep"]["schemaVersion"] = 4 + with self.assertRaisesRegex(RuntimeError, "schemaVersion is 4, expected 5"): + validate_raw_outputs( + make_args(), + outputs["throughputSweep"], + outputs["schedulerPrefill"], + outputs["arrivalInvariance"], + ) + + def test_explicit_config_is_forwarded_to_every_phase_prefix(self): + args = make_args(config="/tmp/gemma-off.toml") + prefix = runner.benchmark_argv(Path("/tmp/darkbloom"), args) + self.assertEqual( + prefix, + [ + "/tmp/darkbloom", + "benchmark", + "--config", + "/tmp/gemma-off.toml", + "--model", + fixtures.MODEL_ID, + ], + ) + for make_command in (runner.sweep_argv, runner.scheduler_argv, runner.arrival_argv): + self.assertEqual(make_command(prefix, args)[: len(prefix)], prefix) + + def test_relative_config_is_normalized_from_repo_root(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + config = root / "configs/off.toml" + config.parent.mkdir() + config.write_text("[gemma_optimizations]\nweighted_r1 = false\n") + self.assertEqual( + runner.resolve_config_path("configs/off.toml", root), str(config.resolve()) + ) + + def test_missing_explicit_config_is_refused(self): + with tempfile.TemporaryDirectory() as directory: + with self.assertRaisesRegex(RuntimeError, "not readable"): + runner.resolve_config_path("missing.toml", Path(directory)) + + def test_baseline_with_different_effective_settings_is_refused(self): + args = make_args() + configuration = baseline_configuration(args) + configuration["gemmaOptimizations"] = opposite_gemma_optimizations() + baseline = {"configuration": configuration} + with self.assertRaisesRegex(RuntimeError, "gemmaOptimizations"): + validate_configuration_pin(args, baseline, fixtures.GEMMA_OPTIMIZATIONS) + + def test_gemma_axis_requires_different_settings(self): + args = make_args(comparison_axis="gemma-optimizations") + baseline = {"configuration": baseline_configuration(args)} + with self.assertRaisesRegex(RuntimeError, "requires different"): + validate_configuration_pin( + args, + baseline, + fixtures.GEMMA_OPTIMIZATIONS, + comparison_axis=args.comparison_axis, + ) + + validate_configuration_pin( + args, + baseline, + opposite_gemma_optimizations(), + comparison_axis=args.comparison_axis, + ) + + def test_gemma_axis_refuses_absent_baseline_provenance(self): + args = make_args(comparison_axis="gemma-optimizations") + configuration = baseline_configuration(args) + del configuration["gemmaOptimizations"] + with self.assertRaisesRegex( + RuntimeError, "baseline configuration did not report effective" + ): + validate_configuration_pin( + args, + {"configuration": configuration}, + opposite_gemma_optimizations(), + comparison_axis=args.comparison_axis, + ) + + def test_gemma_axis_refuses_malformed_baseline_provenance(self): + args = make_args(comparison_axis="gemma-optimizations") + configuration = baseline_configuration(args) + configuration["gemmaOptimizations"]["weightedR1"] = "on" + with self.assertRaisesRegex( + RuntimeError, "baseline configuration reported malformed" + ): + validate_configuration_pin( + args, + {"configuration": configuration}, + opposite_gemma_optimizations(), + comparison_axis=args.comparison_axis, + ) + + def test_gemma_axis_refuses_partial_baseline_provenance(self): + args = make_args(comparison_axis="gemma-optimizations") + configuration = baseline_configuration(args) + del configuration["gemmaOptimizations"]["environment"][ + "MLX_GATHER_QMM_EXPERT_SLICES" + ] + with self.assertRaisesRegex(RuntimeError, "environment projection"): + validate_configuration_pin( + args, + {"configuration": configuration}, + opposite_gemma_optimizations(), + comparison_axis=args.comparison_axis, + ) + + def test_baseline_iterations_mismatch_is_refused(self): + args = make_args() + configuration = baseline_configuration(args) + configuration["iterations"] += 1 + with self.assertRaisesRegex(RuntimeError, "iterations"): + validate_configuration_pin( + args, {"configuration": configuration}, fixtures.GEMMA_OPTIMIZATIONS + ) + + def test_baseline_decode_iterations_mismatch_is_refused(self): + args = make_args() + configuration = baseline_configuration(args) + configuration["decodeIterations"] += 1 + with self.assertRaisesRegex(RuntimeError, "decodeIterations"): + validate_configuration_pin( + args, {"configuration": configuration}, fixtures.GEMMA_OPTIMIZATIONS + ) + + def test_gemma_axis_requires_identical_binary_and_metallib(self): + baseline = { + "metadata": {"binarySha256": "binary-a", "metallibSha256": "metal"} + } + with self.assertRaisesRegex(RuntimeError, "identical artifacts"): + validate_artifact_pin( + baseline, {"binarySha256": "binary-b", "metallibSha256": "metal"} + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/gemma_contbatch/tests/test_kv_backend.py b/scripts/gemma_contbatch/tests/test_kv_backend.py index cb1667a63..f93fbb298 100644 --- a/scripts/gemma_contbatch/tests/test_kv_backend.py +++ b/scripts/gemma_contbatch/tests/test_kv_backend.py @@ -40,6 +40,8 @@ def make_args(**overrides) -> argparse.Namespace: values = { "model": fixtures.MODEL_ID, + "config": None, + "comparison_axis": "code", "iterations": ITERATIONS, "prefill_lengths": list(PREFILL_LENGTHS), "batch_sizes": list(BATCH_SIZES), @@ -208,14 +210,14 @@ def test_a_deliberate_auto_run_forwards_auto_everywhere(self): class DefaultPostureTests(unittest.TestCase): - def test_defaults_are_paged_and_reach_eight(self): + def test_defaults_are_contiguous_and_reach_eight(self): with mock.patch.object(sys, "argv", ["benchmark-gemma-contbatch.py"]): args = parse_args() - self.assertEqual(args.kv_backend, "paged") + self.assertEqual(args.kv_backend, "contiguous") self.assertEqual(args.kv_backend, DEFAULT_KV_BACKEND) self.assertEqual(args.batch_sizes, list(DEFAULT_BATCH_SIZES)) - # The claim this release makes (1.17x aggregate) lives at B=8; a curve - # that stops earlier cannot observe it. + # B=8 remains a useful stress point even though stock serving caps at + # B=4 under the contiguous release posture. self.assertEqual(max(args.batch_sizes), 8) def test_auto_is_still_reachable_for_a_deliberate_run(self): diff --git a/scripts/gemma_contbatch/validation.py b/scripts/gemma_contbatch/validation.py index 7cde4c024..d7c756fcf 100644 --- a/scripts/gemma_contbatch/validation.py +++ b/scripts/gemma_contbatch/validation.py @@ -9,6 +9,13 @@ from .checks import assert_finite, require_positive +RAW_SCHEMA_VERSIONS = { + "throughput sweep": 5, + "scheduler prefill": 2, + "arrival invariance": 4, +} + + def validate_prefill(args: argparse.Namespace, sweep: dict) -> None: expected_prefill_counts = Counter( {length: args.iterations for length in args.prefill_lengths} @@ -102,6 +109,12 @@ def validate_raw_outputs( ("scheduler prefill", scheduler), ("arrival invariance", arrival), ): + expected_schema = RAW_SCHEMA_VERSIONS[name] + if payload.get("schemaVersion") != expected_schema: + raise RuntimeError( + f"{name} schemaVersion is {payload.get('schemaVersion')!r}, " + f"expected {expected_schema}" + ) if payload.get("modelID", "").replace("\\/", "/") != args.model: raise RuntimeError(f"{name} returned the wrong model ID") assert_finite(payload, name)