perf(mlx-swift): R1 expert-QMM mirrors, GemmaQMM diagnostics API, compiled-fn lock-order fix - #11
Conversation
… compiled-fn lock order
- Regenerate Source/Cmlx/mlx-generated quantized.{h,cpp} byte-matched to
the nested canonical R1 source (expert tiles + parallel builder); the
corresponding dispatch mirror is the only regeneration.
- Expose the Apple-side diagnostics contract
(mlx_metal_gemma4_expert_qmm_diagnostics{,_snapshot,_reset}) as Swift
GPU.Gemma4ExpertQMMDiagnostics / gemma4ExpertQMMDiagnostics() /
resetGemma4ExpertQMMDiagnostics() via new transitive headers; no mlx-c
gitlink advance.
- Transforms+Compile: take evalLock before the per-function lock so an
outer compile trace can call a nested compiled function while another
thread evaluates it; fixes the deadlock found in the v0.8.2
verification wave.
- Add SortedGatherQuantizedMMTests: exact-shape ON/OFF parity, selector
misses, NAX priority, AOT absence, and counter invariant/reset.
- Advance nested Cmlx/mlx gitlink to the R1 kernel commit.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: db9b77da44
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
mlx-api.h only defines MLX_API under __cplusplus; the framework umbrella is also parsed in C mode (Swift/ObjC consumers of the Cmlx Clang module), where the C-facing extern-C declarations then failed with 'unknown type name MLX_API', breaking the Xcode framework build. Fall back to an empty definition when the macro is absent. Verified: clang -x c and -x c++ syntax checks on the header, and a negative control that the pre-fix header reproduces the C-mode parse failure. Addresses Codex review on #11.
|
Codex review addressed in P1 — Verification: @codex review |
…cite canonical code in beta docs; bump pins - BenchmarkCommand now loads the runtime snapshot, then routes through the same Start.prepareServeRuntime seam as the serve path: the provider.toml Gemma projection applies before GPUEnforcement's first MLX device access, so a rollback A/B benchmark can no longer measure the default-enabled stack against an operator-disabled config (Codex P1 on #607). - docs/provider/beta-features.md gains canonical file:line references for every behavioral claim (docs/AGENTS.md code-wins rule): settings defaults, startup projection seam, projection/apply authority, the coupled weighted+R1 control (Codex docs P1). - Pins libs/mlx-swift-lm 1d46abc (four shared-tower semantics fixes + contract suite) and libs/mlx-swift cdb7bcd (C-mode MLX_API fallback for the R1 facade; Codex P1 on Layr-Labs/mlx-swift#11). make provider-test: 1666 swift-testing + 82 XCTest pass on this tree.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cdb7bcd807
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…tor to keep the C-mode MLX_API fallback, document ON-path test runs Mirrors the nested R1 review fixes (Source/Cmlx/mlx gitlink advances to 8d538a08): - S1: refresh Source/Cmlx/mlx-generated/metal/quantized.h and the embedded canonical region of Source/Cmlx/mlx-generated/quantized.cpp from the edited canonical kernels/quantized.h; both byte-validated against canonical (the JIT cpp embed = canonical + one blank line, matching the cmake metal-jit codegen output). - S2: mlx-backend-common-gemma4_expert_qmm.h re-syncs the atomic<bool> armed_ counter change through tools/update-mlx-xcodeproj.sh; the script's wholesale regeneration of this header now re-attaches the hand-maintained C-mode MLX_API fallback block after the sed pipeline (guarded on absence, idempotent both directions) instead of wiping it. Verbatim script-section regeneration byte-matches the checked-in header. - S3: Source/Cmlx/include/mlx/gemma4_expert_qmm.h gains _Static_assert ABI pins (sizeof == 80; offsetof armed == 3, attempts == 8, hits == 16, fallback_metallib_unavailable == 72), valid in both C11 and C++17; negative-compiled to prove they fire on drift. - S4: SortedGatherQuantizedMMTests documents the MLX_GATHER_QMM_EXPERT_SLICES=1 ON-path invocation and replaces the vacuous attempts == hits + fallbacks assertion body (attempts() is derived as exactly that sum in C++) with a comment saying why the invariant is constructional; the attempts == 0/1 assertions stay. - S5: Source/MLX/GPU+Metal.swift needs no change - the C diagnostics ABI is untouched by the nested edits (verified against both struct definitions after the sync).
|
Deep-review round complete — fixes landed as
Mirrors re-verified byte-identical to canonical ( @codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…88; generator post-condition Mirrors the nested sortedness fail-safe hardening (Source/Cmlx/mlx gitlink advances to 9b0d1b4c: the adjacent-pair vote in the tile-builder kernel, the fallback_sortedness_retracted counter, and the arm/disarm cycle tests): - S1: refresh Source/Cmlx/mlx-generated/metal/quantized.h (byte-identical to canonical, cmp-verified) and the embedded canonical region of Source/Cmlx/mlx-generated/quantized.cpp (canonical + one trailing blank line; region cmp-verified against canonical). - S2: mlx-backend-common-gemma4_expert_qmm.h re-synced through the real tools/update-mlx-xcodeproj.sh generation; the regenerated header is byte-identical to canonical + the script's sed rewrite rules + the guarded MLX_API attach. Unrelated script regenerations (Cmlx.h whitespace, mlx-c-stream.h doc churn, Metal.hpp one-liner) were reverted -- pre-existing desync, out of scope. Source/Cmlx/include/mlx/gemma4_expert_qmm.h gains the new struct field. - S2b: update-mlx-xcodeproj.sh now enforces a post-condition after the guarded MLX_API attach: the generated header must contain exactly one '#ifndef MLX_API' line, else exit 1. Verified: quiet on the good header, fires on a duplicated guard. - S3: ABI pins updated to the new struct layout: sizeof == 88, offsetof(attempts) == 8, offsetof(hits) == 16, offsetof(fallback_metallib_unavailable) == 72 (unchanged), and offsetof(fallback_sortedness_retracted) == 80. Independently static-asserted against the canonical nested header (C++17); facade negative-compiled with a doctored value to prove the pin fires. - S4: GPU.Gemma4ExpertQMMDiagnostics maps the new counter as fallbackSortednessRetracted (camelCase in the existing fallbackMetallibUnavailable style); the fallbacks sum includes it. Field order/count verified side by side against the C ABI (14 fields). Verified: nested tests 262/262 with 3550 assertions; metal -Wall -Wextra -fno-fast-math zero warnings; fetch-metallib.sh symbol gates (build_gemma4_ sorted_expert_tiles_bm32 + gemma4 tile kernel both present); swift test --filter SortedGatherQuantizedMMTests passes OFF (6 run, 2 flag-gated skips) and MLX_GATHER_QMM_EXPERT_SLICES=1 ON (6 run, 0 skips) against the fresh metallib staged beside the xctest binary; both facades compile in clang -x c (C11) and -x c++ (C++17).
|
Final review wave (#2) complete. Every prior finding is closed; the final verifier pass signed off each repo:
Perf-gate status unchanged: attribution rerun deferred to a cooled machine (see the prior comment — the box measured ~1.7× slow during the late-night window). The remaining open perf item is the nested PR's per-hit @codex review |
|
Codex Review: Didn't find any major issues. Bravo. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Final merge-readiness pass complete:
@codex review |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…onfig controls, source-matched metallib (#607) * chore(deps): bump mlx-swift-lm to query-block prefill attention (#85) abd1985 lands q=128 query sub-blocking in AttentionV1: prompt chunks attend in blocks, slicing K/V to each block's own visible span. Sliding executed/useful collapses from 1.499x to 1.062x at q=128, and the composed-path score tensor becomes O(1) in chunk length (2.04 GB -> 0.51 GB at 124k context on a full-attention layer). Blocked call sites are updateAndAttendRow and borrowAndAttendRow. Decode, last-query prefill, and span-bearing vision chunks keep the single-call path. DARKBLOOM_CBV2_ATTN_QUERY_BLOCK overrides the width; 0 restores the previous behavior. Verified: provider swift build clean; coordinator go build/test clean. provider swift test 1639 tests, 2 pre-existing SelfUpdaterTests failures from a missing local mlx.metallib fixture (environmental, unrelated to this bump). * docs(reports): paged-KV migration plan rev 2, verified against abd1985 Rev 1 was written against submodule b177c35. Re-verified every claim against e65b5bb / abd1985 and corrected 26 items (log in section 22). Four findings changed the plan's shape: 1. #85 already shipped Track A (query sub-blocking), so that track is deleted -- and paged does NOT inherit the win. PagedLayerCache .prefillAttend still issues one SDPA over the full [L, kL] rectangle and adds a gathered K/V copy, so the migration now regresses prefill activation memory unless new item 0.2p lands first. 2. B=8 is NOT reachable by provider config. The coordinator quality cap (concurrency_cap.go:272-299) takes MIN(reported, ceil(1.2 * floor(( solo/15 - 1)/0.27))); gemma-4's 10-18 tok/s solo yields an effective cap of 2 at both B=4 and B=8, pinned by concurrency_cap_test.go :140-147. Gate G0 as written would have measured nothing. Split into G0a/G0b, added section 8 (nine coordinator files, four of them capacity-feedback loops that can eject a B=8 box), tripled Track E. 3. The wire cannot distinguish a paged provider from a contiguous one -- BackendSlotCapacity has no backend discriminator, so a mixed canary is unobservable. Added as Gate G5. 4. CI runs no paged correctness test; four suites are compiled and discarded, and all three named silent-corruption classes have zero fireable assertions. New sections 19 (test/CI) and 20 (benchmark). Also corrected: WS-3.4's justification was false in the dangerous direction (rectangular MTP verification is live in production, so the first paged MTP slot is a process abort, not a fallback); the activation-reserve mirror is servability.go:50, not scheduler.go:1301; the DARKBLOOM_CBV2_PAGED_KV kill switch is unreachable under today's default; Track T should extend CBv2KVSharingParityTests rather than start over. Rev 1's 27,136 donation floor was right and the findings doc's 26,624 was wrong. Also commits the two companion reports the plan cites, so its cross-references resolve in-tree. * perf(provider): Gemma 4 26B v0.8.2 — retained optimization cutover, config-backed controls, source-matched metallib Production integration of the retained v0.8.2 Gemma 4 slices; pins libs/mlx-swift-lm (53e8ccd) and libs/mlx-swift (db9b77d, nested 9dd10bef). - Shared-tower cutover: production/benchmark resolution unwraps the outer MLXVLM.Gemma4 once and threads its owned textModel through engine, cache, sizing, media, prefix/frozen-replay, and MTP paths. EngineV2VLMTextExtraction and its parity-probe plumbing are deleted. - Config authority: GemmaOptimizationEnvironment/Settings resolve the retained layer-18 and coupled weighted-unsort + safe-R1 requests once per process before any MLX device access (launchd/foreground/local/ benchmark), with requested/effective reporting; raw env sampling stays experimental-only. - Version surfaces: ProviderCore.version 0.8.2; coordinator LatestProviderVersion fallback 0.8.2 in sync. - Source-matched metallib: fetch-metallib builds from libs/mlx-swift/Source/Cmlx/mlx (JIT off), refuses artifacts missing NAX/GEMV/the R1 builder/tile symbols; ci/integration/release workflows replace PyPI metallibs and purge cache-restored ones; Makefile provider-build/test stages the metallib for runner and xctest bundles; PackagedRuntimeSmoke/darkbloom runtime-smoke cover it. - e2e/testbed + scripts/gemma_contbatch: benchmark harness environment threading; docs (beta-features, cli-reference, installation, quickstart) and READMEs updated. Measured effect (gemma-4-26B-A4B-it-qat-4bit, M4 Max, immutable pre-edit vs retained weighted+R1 profile, median/geomean): arrival TTFT +44.8%, arrival end-to-end TPS +19.6%, single-shot TTFT +15.2% (2048-tok 1715ms -> ~1389ms), decode aggregate TPS +1.3%, prefill TPS +2.4%, peak memory unchanged. Dropped: expert gate/up packing, dense gate/up packing, standalone weighted-unsort, standalone R1, historical accepted R1 control (noisy/negative). R1 kernel armed but recorded zero production hits; retained opt-in under the coupled gate. Benchmark artifacts predate the final kernel-body source edit; rebuild and re-verify before release claims. * docs(changelog): v0.8.2 candidate entry for the Gemma 4 optimization stack * test(provider): use closure form in GemmaOptimizationReportingTests #expect allSatisfy(\.effective) trips the macro's throwing-call detection; closure form is identical and compiles. * fix(provider): apply Gemma config before first benchmark MLX access; cite canonical code in beta docs; bump pins - BenchmarkCommand now loads the runtime snapshot, then routes through the same Start.prepareServeRuntime seam as the serve path: the provider.toml Gemma projection applies before GPUEnforcement's first MLX device access, so a rollback A/B benchmark can no longer measure the default-enabled stack against an operator-disabled config (Codex P1 on #607). - docs/provider/beta-features.md gains canonical file:line references for every behavioral claim (docs/AGENTS.md code-wins rule): settings defaults, startup projection seam, projection/apply authority, the coupled weighted+R1 control (Codex docs P1). - Pins libs/mlx-swift-lm 1d46abc (four shared-tower semantics fixes + contract suite) and libs/mlx-swift cdb7bcd (C-mode MLX_API fallback for the R1 facade; Codex P1 on Layr-Labs/mlx-swift#11). make provider-test: 1666 swift-testing + 82 XCTest pass on this tree. * fix(provider): harden config loading surface + benchmark/CLI integrity follow-ups Review follow-ups on the v0.8.2 Gemma 4 optimization stack (#607): - R1 (HIGH, silent rollback defeat): the lenient ConfigManager.parse (whole-config defaults on ANY decode error) stays only as the test-facing API; production file loads route through the new strict ConfigManager.parseValidating via ConfigManager.load(from:), throwing ConfigError.parseFailed with the decode detail. A malformed [gemma_optimizations] entry (weighted_r1 = 0, prefill_layer18 = "false") previously re-armed the whole default-on stack with zero log; every snapshot-based command (start/benchmark/beta/status/...) now refuses at load. Missing files keep defaulting one layer up (loadDefault/loadRuntimeSnapshot), and missing sections/keys keep their per-key default-on decode. - R2 (HIGH, benchmark A/B integrity): darkbloom benchmark now reads the three low-level env keys before applying the projection and FAILS (stderr naming key, shell value, config value + ExitCode.failure) when a shell preset conflicts, instead of silently overwriting it and invalidating the runner.py os.environ artifact metadata; on success it prints one line with the effective controls. Benchmark A/B is explicitly config-driven (docs/provider/beta-features.md updated with the new citations/line numbers). - R3 (MED, seam placement): the ordering seam moved from Start.prepareServeRuntime into the new ServeRuntimePreparer.prepareRuntime (config projection strictly BEFORE the first MLX touch; a rejected projection aborts before engine construction). Start.prepareServeRuntime remains as a forwarding shim for compatibility (StartCommand.run + one shim test); BenchmarkCommand and StartCommandTests now target ServeRuntimePreparer directly. - R4 (LOW): darkbloom beta enable/disable no longer no-ops when the key (or its [section]) is ABSENT but the decode default matches the target: the value is materially written (BetaFeature.configAddress + tomlKeyPresent), pinning operator intent against future default flips. A key already pinned at the target is still a true no-op (no rewrite). - R5 (LOW): CHANGELOG stat fixes vs the committed-tree tree benchmark artifact: arrival TTFT speedups 1.411/1.385/1.470/1.762 -> "1.38-1.76x" (was "1.39-1.76x"); "peak memory unchanged" replaced with the retention memory control numbers (B1/B2/B4 peaks 14.0/14.64/15.05 GiB vs default 14.06/14.70/15.11; tmp/benchmarks/v0-8-2-retention-comparison.json). - R6 (LOW): the beta load->modify->save window is serialized by an exclusive flock(2) on a stable provider.toml.lock SIDECAR file, with the config re-loaded inside the lock (lost-update RMW race). Sidecar chosen over locking the config file itself: ConfigManager.save writes atomically via temp-file+rename, so the config inode changes on every save and concurrent writers holding different inodes would NOT exclude each other; the sidecar path is never renamed, giving every contender the same inode. fd close releases the lock even across a throw. R7 (note only, semantics unchanged): GemmaOptimizationEnvironment holds no armed/latched global state -- projection() is pure and apply() only setenv()s the three latches at process start. The POSIX process environment is global and getenv/setenv are not thread-safe, which stays correct because both call sites (start, benchmark) apply before MLX init and before any worker concurrency; a future call site that applies mid-run would need its own serialization. Tests: 1692 swift-testing tests + 82 XCTest, all passing (make provider-test). New: ConfigValidationTests (11), BetaCommandTests (10), RuntimeSnapshotConfigTests (2), StartCommandTests +4 (env-conflict guard + shim forwarding; pre-existing Start seam tests retargeted). * chore(deps): pin mlx-swift-lm f00c9bd + mlx-swift 746d7c5 (review-fix cascade); fix benchmark stdout hygiene - mlx-swift-lm f00c9bd: completes the global-KV-head rule across the CBv2 layer-kind derivation and MTP validation, re-forges the vacuous fp16 overflow proof at norm gains, normalizes short/empty layer_types, pins rank-1 at both entries and the canonical ProportionalRoPE construction. - mlx-swift 746d7c5 (with nested 8d538a08): fail-safe sortedness check in the R1 tile builder (host re-routes to legacy on violation), shared route predicate, hit-only bias normalization, atomic armed_; mirrors/facades byte-synced; xcodeproj generator now preserves the C-mode MLX_API fallback; facade ABI pinned with _Static_asserts; ON-path test invocation documented. - SlotSizingDriftTests: track the gemma4LayerKinds signature change. - BenchmarkCommand: project the effective settings echo to stderr — benchmark subcommands emit machine-parsed JSON on stdout (a stdout line broke the gemma-contbatch harness). * fix(provider): scope UpdateCommand config fallback to missing files; widen preload-gate margin under parallel load; bump pins - UpdateCommand: only ConfigError.readFailed falls back to defaults — an existing-but-malformed provider.toml now fails loudly instead of silently resetting every setting and aiming the update check at the wrong coordinator (final root-verifier residual item; watchdog's deliberate fail-open posture unchanged). - StartupPreloadTests: wall-clock margin 2.8s -> 4.5s; the semantic signal is the .timedOut outcome + background continuation, and the tight margin flaked at 2.8-3.5s under full-suite parallelism on loaded machines (1692-test runs, twice, on this box). - Pins: libs/mlx-swift-lm 06046c3 (last k_eq_v-gated head rule gone from validateAttentionProducts), libs/mlx-swift 38eaa93 (nested 9b0d1b4c: sound adjacent-pair sortedness oracle + dedicated fallback_sortedness_retracted counter, ABI 88, generator post- condition; mirrors/facades byte-synced). * chore(deps): pin upstream-integrated mlx-swift-lm * chore(deps): pin final Gemma review fixes * test(provider): synchronize metallib environment mutations * fix(bench): harden Gemma baseline verification * docs(provider): align preflight and KV citations * fix(bench): enforce fixed decode token budgets * chore(deps): pin CBv2-scoped weighted expert path * fix(ci): remove unused schema helper * chore(deps): pin merged Gemma safety fixes * fix(bench): default Gemma runs to contiguous KV * fix(provider): reject all-mode Gemma CBv2 * fix(provider): preserve update config read failures * chore(deps): pin final Gemma compatibility fixes * chore(deps): pin merged mlx-swift-lm main
What
Swift-side integration for the Gemma 4 expert-QMM (R1) kernel, plus one concurrency fix found during verification:
Source/Cmlx/mlx-generated/metal/quantized.handSource/Cmlx/mlx-generated/quantized.cppregenerated from the nested canonical source (perf(mlx): opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder mlx#4); the generated header is verified byte-identical to its canonical source and only the corresponding dispatch mirror was regenerated.include/mlx/gemma4_expert_qmm.h, framework umbrella entries) exposingmlx_metal_gemma4_expert_qmm_diagnostics{,_snapshot,_reset}, wrapped in Swift asGPU.Gemma4ExpertQMMDiagnostics/GPU.gemma4ExpertQMMDiagnostics()/GPU.resetGemma4ExpertQMMDiagnostics()with derivedfallbacks.CompiledFunction.callnow takesevalLockbefore the per-function lock, so an outer compile trace can call a nested compiled function while another thread is evaluating it. This resolved a real deadlock hit in the v0.8.2 verification wave (nested compiled decode under continuous batching).Tests/MLXTests/SortedGatherQuantizedMMTests.swift: exact-shape ON/OFF arithmetic parity, one-field selector misses, NAX priority, AOT absence, counter invariants/reset.Source/Cmlx/mlxgitlink to the R1 kernel commit (perf(mlx): opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder mlx#4).Before / After — behavior
flowchart LR subgraph Before A1[Swift caller] --> B1[no way to observe<br/>expert-QMM engagement] C1[outer compile trace calling<br/>nested compiled fn] --> D1[possible deadlock:<br/>fn-lock → evalLock vs evalLock → fn-lock] end subgraph After A2[Swift caller] --> B2[GPU.gemma4ExpertQMMDiagnostics:<br/>requested / aotAvailable / naxAvailable /<br/>hits + per-class fallbacks; reset()] C2[outer compile trace calling<br/>nested compiled fn] --> D2[one global order:<br/>evalLock → fn-lock everywhere] endBefore / After — code
flowchart TD subgraph Before M1[mlx-generated mirrors<br/>pre-R1] --> M2[no gemma4 headers] T1[Transforms+Compile.call:<br/>per-function lock only] --> T2[trace path acquires<br/>evalLock after fn lock] end subgraph After N1[nested Cmlx/mlx gitlink<br/>Layr-Labs/mlx#4] --> N2[mlx-generated quantized.h/.cpp<br/>byte-matched mirrors] N2 --> N3[include/mlx/gemma4_expert_qmm.h<br/>+ framework umbrella] N3 --> N4[GPU+Metal.swift diagnostics] U1[Transforms+Compile.call:<br/>evalLock → per-function lock] --> U2[trace path keeps same order,<br/>no re-entrant inversion] endTesting
swift teston the downstream tree (1666 tests / 178 suites) passes with these changes linked.3f9d85ec…) to the artifact measured in the downstream production retention matrix.Stack: sits on Layr-Labs/mlx#4; downstream of Layr-Labs/mlx-swift-lm → d-inference v0.8.2.