Skip to content

Cluster collectives, quantized numerical fixes and optional Metal allocation cap - #35

Merged
davidtai merged 53 commits into
mainfrom
darkbloom/sealed-all-sum
Oct 28, 2026
Merged

davidtai merged 53 commits into
mainfrom
darkbloom/sealed-all-sum

Conversation

@jonathan308

@jonathan308 jonathan308 commented Oct 15, 2026 •

Copy link
Copy Markdown

Progress 2026-10-22 09:00 PDT: sealed reduce-scatter and all-gather

reduce_scatter_sealed and all_gather_sealed for two ranks seal and authenticate every piece through the same codec as all_sum_sealed. Every sealed collective now draws its counter from one process-wide sequence; a graph of all-sums only gets the same values as before.

On the two-Mac pair over JACCL:

  • Results: reduce-scatter equals this rank's rows of the sealed all-sum, and all-gather equals the padded all-sum, bit for bit. Checked at five shapes (f32, bf16, f16, up to 8192 x 4096).
  • Tamper: a flipped byte in a peer's piece is refused on both ranks for every kind: all-gather at counters 0 and 200, reduce-scatter at 2 and 202, all-sum at 1, 3, 201 and 203.

The first user, GLM sequence-parallel prefill, is bitwise but slower on the pair, so it stays off; the collectives are kept for it and later uses.

Progress 2026-10-21: rebased onto the latest main/master, candidate 24

Head 79d2ebe8.

  • Merged the single-Mac core (Single-Mac GLM-5.3-Flash native fixes without JACCL #37, now on main as ad4745a2). Its general commits are cherry-picks of this branch's, so the merge keeps this branch's tree, adds main's [Metal] thread-scoped command-buffer commit limits #36 (thread-scoped command-buffer limits), and records main's squash.
  • Reverted 4d913204 (page-backed footprints) explicitly, as the single-Mac core did. It charged GLM MTP requests about 11.3 GB instead of 4.1 GB.
  • fork.yaml now describes the host cuts and the streamed sealed sum; the fork-diff gate passes with 233 of 233 changed files described.
  • The merged core's metallib is byte-identical to candidate 20's (79567e78).

Verification (two M5 Ultra 256 GB Macs).

  • Pair test tool: tokens and top logits identical to the previous verified build, candidate 20, on the 4K plain, 4K MTP-depth-2 and 16K reference runs.

  • Cluster check suites: 30 of 37 pass. Six failures were already present on candidate 20 (stale check-runner file lists, the GPT-OSS stage check, a 32K qualification test). PlacementChecks needs its tests updated for the memory change.

  • Endpoint (candidate 24, tensor parallel, chosen by the planner):

    Measure Candidate 24 Candidate 20
    Prefill, 2.7K / 10.5K / 31.5K tokens (tok/s) 2,966 / 3,498 / 3,526 2,951 / 3,511 / 3,512
    Decode, 256 / 1,024 tokens (tok/s) 87.4 / 88.1 (92.1 / 86.2 over 5 reps) 87.9 / 85.5

    Prefix reuse returns the first token in 0.87–1.25 s. Smoke checks pass (391, stop string, length).

Known gap, being fixed next: the cluster serves one request at a time. Aggregate throughput stays at about 80 tok/s from 1 to 16 concurrent streams, and the other streams queue. Concurrent serving with real batching, for every model in both TP and pipeline, is the next change on this stack.


Single-Mac split, 2026-10-20

The general (non-JACCL) commits of this branch are also on Layr-Labs/mlx#37 (work/glm53-flash-release → main), cherry-picked with their (cherry picked from …) trailers. Single-Mac GLM-5.3-Flash (d-inference #1516) now pins that chain instead of this one, so it can merge first. This branch is unchanged. Once the single-Mac PR merges, this stack merges main (fork.yaml needs a manual resolution) and the two-Mac pair is re-verified bitwise.


Summary

Adds sealed two-rank all-sum and repairs quantized Metal kernels used by the cluster engine. Each partial result crosses the link in a frame sealed by the caller's codec. Sorted NAX gathers bound partial activation tiles, and legacy/NAX affine/FP loaders have distinct symbols and bounded weight reads. Scalar CPU quantized matmul now accumulates in float instead of rounding every partial sum to BF16/FP16. An optional install-once Metal backing cap reserves pending bytes before allocation and returns credit only after Metal releases and unmaps the backing; residency-registration failure also cleans up ordinary allocations.

Capped no-copy backing is prefaulted before Metal receives it. JACCL adds a one-shot bounded startup wait, with refusal and failure cleanup preserved. Custom Metal kernels validate the declared launch geometry and may retry pipeline compilation with an explicit 1024-thread contract; tensor math and dispatch dimensions are unchanged.

The detached Metal footprint policy includes a VM-page minimum, so a capped 384-byte receive with 16,384 bytes of backing fits its conservative 32,767-byte bound. Actual backing allocation, cap accounting and ownership checks are unchanged.

Stacked on #33, with its current target integrated. The Swift and provider consumers must repin this revision; their release and real-model RDMA qualification remain separate.

Before

flowchart LR
  T["Two-rank partial results"] --> P["Plain collective payload"]
  G["Sorted quantized gather"] --> K["NAX tail reads beyond remaining K"] --> B["Wrong values or NaNs"]
  A["Offline Metal link"] --> L["Different loaders share a symbol"] --> B
  Q["CPU quantized product"] --> R["Low-precision partial sums"] --> B
  M["Soft memory limit"] --> U["Backing can exceed an experiment allowance"]
  F["384-byte capped receive"] --> B0["Detached bound: 767 bytes"] --> X["Valid page-backed receive rejected"]
  J["Startup collectives"] --> W["Ordinary idle timeout during load"]
  H["Required 1024-thread kernel"] --> R0["Compiler-selected pipeline capacity too small"]
  classDef default fill:#fff1f2,stroke:#be123c,color:#881337
Loading

After

flowchart LR
  T["Two-rank partial results"] --> C["Caller codec"] --> S["Sealed all-sum frames"]
  G["Sorted quantized gather"] --> K["Bound rows, packed columns and remaining K"] --> V["Exact valid outputs"]
  A["Offline Metal link"] --> L["Distinct loader and helper identities"] --> V
  Q["CPU quantized product"] --> R["Float accumulation; one output cast"] --> V
  M["Opt-in cap"] --> P["Reserve backing before allocation"] --> D["Refund after actual deallocation"]
  F["384-byte capped receive"] --> B0["Detached bound: 32767 bytes"] --> X["Backing covered; ownership checks retained"]
  J["Startup collectives"] --> W["Fixed one-shot deadline then ordinary idle policy"]
  H["Validated 1024-thread kernel"] --> R0["Bounded pipeline retry; same geometry and math"]
  classDef default fill:#ecfdf5,stroke:#047857,color:#064e3b
Loading

Evidence

Current successor qualification

Current signed head is 54975cd; C #15 and Swift #56 now select this revision. The actual target base remains 56a2043.

  • Before: the actual custom-kernel evaluator regression rejects required full-size launches when the initial Metal pipeline exposes insufficient capacity. After: 33 evaluator checks pass under UBSan; invalid geometry still refuses. On the M3 Ultra, 17 actual Metal bounds/dispatch checks and all 21 GLM dtype/geometry cases pass with the required threadgroup size unchanged.
  • Sixteen separate fresh-process allocator/cap modes pass, including physical footprint growth, retained backing, concurrency, foreign buffers, GPU evaluation and registration failure. The footprint assertion now rejects a decreasing counter before unsigned subtraction; five arithmetic controls cover that regression.
  • Startup policy: 23 ASan/UBSan controls and seven real mesh/ring dispatch controls with simulated verbs pass. The actual C bridge passes 17 refusal cases plus valid forwarding. Simulated verbs are not real RDMA proof.
  • The source-matched SDK composition passes 16 focused stage/native/MTP-fixture tests, including the new prefill readout and cache/failed-ack behavior. This is component evidence, not clustered MTP support or whole-model performance qualification.
  • Configured hooks pass; actual-HEAD fork coverage is 225/225. Independent source/spec/refactor review found no remaining scoped defect; public core production bytes match the tested canonical nested source.
  • Current hosted lint, fork, ASan, UBSan and both Linux Fedora jobs pass; Metal CI is still running at the latest readback. Swift consumer CI has a build failure under investigation. Full current consumer suites, final installed-pair qualification and applicable modalities remain open. No release-readiness claim.

Historical page-backed subpage evidence

  • Before: the real worker receive regression fails four assertions at 256/384 bytes: actual capped backing is 16,384 bytes, exceeding bounds 511/767; the production owned-receive validator refuses it. The standalone real capped allocator independently reproduces the 256-byte mismatch. A missing-metallib setup failure is recorded separately from those functional failures.
  • After: the same worker regression passes (one test, zero failures/skips), including real shared/offset-view refusal. Fourteen fresh-process native allocation/cap controls pass: ordinary/capped footprint, install, retained, churn, foreign, GPU, late, nil-foreign, ordinary/capped registration, foreign-registration failure and concurrent reservations. The formatted public standalone test was freshly compiled and is byte-identical to that tested executable, with the same native object and AOT library; no duplicate GPU run is claimed.
  • Both installed ranks subsequently complete the authenticated encrypted GLM data-plane exchange over JACCL. Loading still refuses under file-cache/compression pressure. This establishes the receive-bound correction, not full-model inference, multimodal serving, endpoint readiness or a serving grant.
  • All five configured full-file pre-commit hooks pass; fork coverage is 214/214 at the actual upstream merge-base. Independent SOURCE/SPEC/REFACTOR review found no actionable findings and justified no refactor. Swift consumers must repin the public successor.
  • Historical CI observation: preceding head 9e22e1998a3b6f7f6f4ad58dc2038ecf9f4d0cb9 failed CPU test_double.TestDouble.test_gemm at its float32/float64 allclose assertion. The hosted log is now available; its cause remains unresolved. See the current successor CI disposition above; no tolerance relaxation or skipped check is included.

The following larger suites describe earlier compositions and unchanged component coverage; they are not a fresh full-suite run of this successor.

  • Earlier combined cap/CPU composition, M3: 12 fresh-process Metal cap cases and 1 Swift cap case pass; full native Swift/Cmlx 931 passed, 2 expected skips, 0 failures (74.75 s). Source, native binary, embedded JIT and AOT identities were unchanged during qualification. The skipped cap case passes separately with its required fresh-process opt-in; the other skip is the opt-in performance test.
  • CPU numerical regression: BF16 produced 592 instead of the exact 576. Exact-value regressions reproduce six affine and fifteen adjacent mode/orientation subcase failures, then pass with float accumulators. Full local Debug/coverage CPU workflow: 292 C++ passed; 815 Python passed, 62 declared skips, 0 failures across 33 modules. Python 3.14 and the local compiler differ from hosted CI; this is a confirmed local defect, separate from the now-accessible hosted CPU GEMM failure above.
  • Cap failure and concurrency controls: retained buffers stay charged, pending real-Metal reservations cannot exceed the cap, cache transfers preserve credit, failed registration rolls back, and failed foreign wrappers cannot silently copy. M5 cap production passed 12 cases plus its full native suite before the CPU-only delta; this is distinct earlier-composition evidence.
  • At that earlier composition, pinned hooks and core fork coverage passed at 213/213. Independent reviews cover allocator lifetime, CPU arithmetic, failure cleanup and concurrency. Core, C and Swift consumer pins must move together.

Earlier kernel/execution evidence below remains applicable to unchanged kernel sources; it is not an exact current M5 combined-binary claim.

  • Tail regression: the old Swift composition failed 48 of 144 geometry cases. The repaired regression passes 272 cases across affine 2/3/4/5/6/8-bit, MXFP4/MXFP8/NVFP4, packing groups, FP32/FP16/BF16, partial tiles and adjacent NaNs.
  • AOT regression: the baseline fails 32 subcases; the tail-only repair still fails 16 MXFP4 subcases. Renaming only the FP NAX loader changes NaNs into the exact expected outputs with an unchanged Python module. The final separation of all incompatible loaders passes the 128-case Python regression and all 39 quantization tests, using the real precompiled Metal library on M5 Max.
  • Legacy M3 regression: published core 1fb86d3 fails 26 exact NVFP4 partial-K assertions. This successor fixes legacy row/column bounds; the same focused tests pass without assertion changes. Final M3 native suite: 931 passed, one opt-in performance skip, zero failures.
  • Swift integration: the final source-matched M5 native suite passes 930 tests, with 2 explicit skips (opt-in performance and non-NAX-only retract coverage), zero failures. The final AOT library passes all 39 Python quantization tests on both M3 and M5; the targeted 128-cell regression also passes on M3. Expert-tile kernels are enabled and the resource-count limit is 50,000. All four generated kernel strings are verified present byte-for-byte in the tested executable. Broader SDK qualification continues in the consumer PR.
  • Configured pre-commit checks pass. Fork-diff coverage: 209/209 changed files described at that earlier head against the actual upstream merge-base. Independent review confirmed tail bounds and identifier-only loader separation; no dispatch disablement, precision change or tolerance relaxation.
  • Existing sealed all-sum hardware evidence remains the two-Mac runs linked from d-inference #1407. These new component results do not establish GLM RDMA or installed-endpoint readiness.

Merge Danger

Door: two-way. No storage migration or deployment.

Blast Radius: native. Prefaulting changes when capped pages become resident, and the new startup deadline and bounded pipeline retry affect native execution. The conservative detached minimum changes quotations for subpage buffers; it does not change physical allocations or cap enforcement. Quantized Metal loader identities and partial-tile reads change; consumers must refresh both their core pin and embedded JIT sources. CPU scalar accumulation precision also changes. The cap limits owned Metal backing and conservatively charged foreign wrappers, not whole-process/host/driver memory; host budget and OS headroom remain separately required. Cap installation is opt-in, once-only, and must precede positive buffer attempts. Formatting-only updates to existing sealed sources satisfy the configured formatter.

Implementation and regression investigation used Codex, with an independent agent review. The original sealed collective was prepared with Claude Code.

512-thread compiler-capacity follow-up

A GLM fused KDA test requested 512 threads on a pipeline compiled with a 384-thread limit. The existing retry covered only 1024-thread groups. The core now retries complete 512- or 1024-wide groups with an explicit compiler-capacity declaration, preserving shader math and dispatch geometry. Device bounds and the selected pipeline’s actual capacity still reject unsupported launches.

  • Before: SDK native CI failed in GLM5NextLinearDecodeKernelTests/decodeStepMatchesComposedStepOnGPU with the 512-versus-384 capacity error.
  • After: the source-matched native consumer suite passed 32 tests across 8 suites, zero failures or skips, including the actual GLM linear-decode kernel suite. The CPU regression using the evaluator’s actual retry block changed from failing to passing across five retry/rejection cases.
  • Normal formatting and fork-description checks pass. Python launch geometry coverage includes 512; that added Python case was not separately executed. These are native correctness results, not an isolated performance benchmark or a claim that all SDK CI is green.

Update 2026-10-20: fast-synch GPU wait ordered after pending passes (f0db8cf), 1 MiB JACCL frames (61fd8a8)

  • Hang fix. A GPU pass that spins on a CPU-stream fence (MLX_METAL_FAST_SYNCH=1) had no dependency on earlier passes still in flight on its queue. Without hazard tracking Metal could start it first, and when an earlier pass was what the awaited CPU work needed (GPU → CPU → GPU, the sealed all-sum between two layers) the GPU spun forever at 100 %. Traced on two M5 Ultra with a fence/command-buffer ring: a committed buffer of plain elementwise ops stayed scheduled, not started while the later spin ran. Fence::wait now marks the spin's pass CommandEncoder::wait_for_all_pending(), which waits on every pending pass's fence of that queue; the CPU work it waits for already follows them.
  • Sealed sum add split across cores (dispatch_apply, 256K-element spans), addends still in rank order (bit-identical). 32 MiB fp32 sealed sum 20–25 ms → 5.9 ms.
  • JACCL largest frame 512 KiB → 1 MiB (BUFFER_SIZES 9), mesh pipeline from MESH_PIPELINE (2, static_assert). 16 MiB bf16 all-sum 3.02 → 2.69 ms. 2 MiB frames and a pipeline of 4 fail in the driver.

Result on GLM-5.3-Flash TP2 (sealed, two M5 Ultra): prefill 1,913 tok/s, decode 50.7 tok/s (MLX_METAL_FAST_SYNCH=0 instead: 13 tok/s decode). Native unsealed TP2 unchanged (2,400 / 56.3).

🤖 Generated with Claude Code

Update 2026-10-20 (evening): host-encode ports (872fa01, 5711fa5), streamed sealed all-sum (8c6f9e7)

All three commits are Verified; nothing is merged.

Host time (ports of the owner's mlx-m5u work):

  • 872fa01 (mlx-m5u 8c4bf63fe): each command buffer retains its primitives' buffers in one completion handler, not one per primitive.
  • 5711fa5 (mlx-m5u 42f2dcefb): a custom kernel's generated source and library name are built once per kernel and template instance.
  • GLM-5.3-Flash TP2 decode on two M5 Ultra: 16.75 -> 15.92 ms a step.

Streamed sealed all-sum (8c6f9e7). On a two-rank JACCL mesh, SealedAllSum now runs one CPU-stream task over a new MeshImpl::exchange:

  • Each 256 KiB sealed piece carries its own AES-256-GCM tag.
  • A piece is sent as soon as it is sealed, and opened and added as soon as it lands.
  • Only the first and last pieces wait for crypto.
  • Ring, TCP and larger groups keep the v1 path unchanged.
flowchart LR
  subgraph Before["v1: seal, gather, open, add"]
    S1[seal whole frame] --> G1[all_gather 2n] --> O1[open whole frame] --> A1[add]
  end
  subgraph After["v2: per piece, pipelined"]
    S2[seal piece k] --> X2[exchange frame] --> O2[open piece k] --> A2[add piece k in cache]
    S2 -. "piece k+1 seals while k crosses" .-> S2
  end
Loading
Property How it holds
Every byte authenticated One AES-GCM box per piece under its direction and generation key. The receiver derives the piece count and lengths from n and the piece size, never from the wire. The sum throws unless every piece opened.
Unique nonces (piece, sender, counter); each piece is sealed once, never resealed.
v1 and v2 kept apart v2 authenticated data = v1 authenticated data plus the piece size, so neither layout's piece verifies as the other's.
Fail closed, no hang A piece that does not open latches the codec. A refused seal sends zeros, so both ranks refuse in lockstep.
Bit-identical results Same per-element expression and rank-0-first operand order as v1.

v1 fix in the same commit. A refused seal now zeroes its frame before it throws. Before, the following gather could send a recycled buffer's stale bytes in the clear.

Measured on two M5 Ultra over Thunderbolt 5 (GLM TP2, sealed bf16 partials, 15.7K-token prompt, 8,192-token chunks):

  • Prefill rose from 3,554 to 3,812 tok/s, with output tokens identical.
  • A tampered piece at sealed counter 0, and again at 200, was refused on both ranks. Both ranks ended without hanging.

tests/jaccl/exchange_tests.cpp uses fake verbs with the real MeshImpl and streamed-sum code. It passes under ASan/UBSan and TSan, covering:

  • sizes from 1 B to 3 MiB;
  • bitwise equality with v1 for fp32, bf16 and fp16 over 1–65 pieces;
  • tamper, seal failure and a closed group.

For platform review. Callers still name the reduction sealed_all_gather_sum_v1 (Qwen pair qualification records, ClusterTensorParallelTerms). The results are bit-identical, so existing qualification records stay valid. The wire layout is bound by the authenticated data and by the build SHA in each load agreement.

🤖 Generated with Claude Code

Update 2026-10-21: host-time cuts and opt-in sum plumbing (6034a78 … 7620956)

Six Verified commits; nothing merged. Every change keeps model outputs bit for bit (GLM-5.3-Flash TP2 on two M5 Ultra: tokens and top logits identical).

Commit Change Default Measured
6034a78 Fused coherent fence update for CPU-stream inputs (MLX_METAL_FUSED_FENCE_UPDATE) off ≤ 0.1 ms/step
a7d4079 Fenced inputs for custom kernels (MLX_METAL_FENCED_INPUTS): the consumer waits on the fence itself, after all pending passes (the fast-synch deadlock ordering is kept) off ≤ 0.1 ms/step
7ebf8e1 Small flat sets for per-dispatch encoder hazard tracking (same barrier/commit/fence decisions, replayed on 100 traces) on ~0.2 ms host/step
9ec3dba Eager commit after a cross-device fence update (MLX_EAGER_FENCE_COMMIT) off neutral
a94584b Separate CPU-stream task budget (MLX_MAX_ACTIVE_CPU_TASKS) off neutral
7620956 Shared generated sources, pipeline resolved once, repeat launches skip name/template/key rebuilds, flat map in eval_impl, cheaper graph teardown on with mlx-c/mlx-swift: one-rank host CPU per decode step 6.5–6.9 → 4.7–4.8 ms

Not compiled here: the CUDA side of the shared-source change (no CUDA toolchain on these Macs).

🤖 Generated with Claude Code

Progress 2026-10-26: main merged, stack re-pinned, validated on two M5 Ultra 256 GB

Head 2b5253ad: #33 and main merged; Check Lint fixed (95938c4b); the streamed-sum exchange test compiles again (2b5253ad).

Nothing merged or released; the owners sign and release.

Update 2026-10-26 17:15 PDT: main/master 0.10.3 merged through the stack

Head 424562b9: the refreshed #33 merged (brings main 6bc9de5b). Every PR in the cluster stack (mlx #33 → #35 → #41, mlx-c #15, mlx-swift ml-explore#54 → ml-explore#56 → ml-explore#61, SDK ml-explore#308 → ml-explore#352, d-inference ml-explore#1407 → ml-explore#1426 / ml-explore#1577, platform ml-explore#54) is again mergeable in that order, as drafts.

Review fixes and known limit (2026-10-27)

Head 169802f4: review fixes in be4b55b1 (fenced custom-kernel inputs stay in the device address space, fenced waits are taken before any launch check can throw, sealed frames over INT_MAX bytes are refused, codec contract documented in sealed.h), #33's fixes merged in, and the clang-format fix 169802f4.

Known limit (for David to rule): if one rank fails before it takes its sealed counter and the other rank does not, the counters desync and every later sealed operation fails authentication. This fails closed, but nothing recovers the group; the group must be rebuilt.

Known limit (for David to rule): MeshGroup::end_startup_wait_budget ends the budget wire by wire. Each wire's deadline is set a moment after the previous one, so at the very edge of the deadline wire 0 can end while wire 1 throws "Cannot end an expired startup budget". The group is then partly ended: wire 1 still enforces the expired deadline on its next call. This existed before the shared failure flag; it was not changed here.

CI fix 94baa025 (tests only, no library change): progress_guard_checks runs with a 2000 ms limit instead of 250 ms (every case scales with the limit, ratio kept), and the startup expiry cases spin their polls. The abort in startup_wait_budget_groups was the fake poll's 8 s safety deadline, not a crash in the library.

…codec-sealed frames)

all_sum_sealed(x, group, stream) is evaluated on the CPU stream inside the
lazy graph like the native all-sum: the local partial is sealed by a codec
the embedding program installs (set_seal_codec: seal, open, context,
per-frame overhead), the two frames are all-gathered over the group, the
peer's frame is opened (a frame that does not authenticate fails the
evaluation; there is no plaintext fallback) and the partials are added in
rank order (rank 0 + rank 1) so both ranks hold the same bits. float32,
bfloat16, float16; a counter advances per sealed all-sum in evaluation
order, identically on both ranks, for the codec's nonces. No codec: the
graph refuses to build.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Use the remaining reduction width for the final activation tile. Check
row and packed-column bounds before reading affine or floating-quantized
weights and scales, then zero padding lanes.

Add exact-value partial-tile and NaN-neighbor regression coverage for
affine, MXFP4, MXFP8 and NVFP4. The companion Swift regression reproduced
48 failing geometries and passes all 272 cases with this repair. The
original MiMo and Qwen fallback failures also pass on M5 Max.

AI assistance: Codex diagnosed, implemented and tested this patch under
the repository owner's explicit direction.
Retain the target branch's explicit fork-diff coverage for the JACCL
memory release hook and completion-queue depth sources. This merge
changes fork.yaml only; the compiled runtime source is unchanged.
@jonathan308 jonathan308 changed the title Sealed all-sum for two-rank tensor parallel (pinned by d-inference#1407) Sealed two-rank all-sum and quantized Metal correctness fixes Oct 18, 2026
@jonathan308 jonathan308 changed the title Sealed two-rank all-sum and quantized Metal correctness fixes Cluster collectives, quantized numerical fixes and optional Metal allocation cap Oct 18, 2026
jonathan308 and others added 14 commits October 18, 2026 11:40
Keep the allocation bound stable before and after process cap installation.
Add real ordinary/capped boundary coverage without changing allocation,
ownership, process-cap enforcement or the policy ABI. Document the bound.

Implemented and validated with OpenAI Codex.
Prefault capped no-copy backing before handing it to Metal. Add a one-shot JACCL startup wait budget with C-facing group operations and failure cleanup tests. Validate custom Metal launch dimensions and bound the 1024-thread pipeline retry to the declared geometry.

CPU policy, bridge, allocation-ledger and evaluator regressions pass. Native GPU and end-to-end qualification are tracked separately.
Raise the largest JACCL frame from 512 KiB to 1 MiB (BUFFER_SIZES 9) and
size the mesh pipeline from MESH_PIPELINE (2), asserted at compile time.
Measured on two M5 Ultra over Thunderbolt 5: a 16 MiB bf16 all-sum drops
from 3.02 ms to 2.69 ms; GLM-5.3-Flash TP2 prefill on a 3.9K-token prompt
goes from 2,130 to 2,289 tok/s. 2 MiB frames and a pipeline of 4 fail in
the driver.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…lel sealed add

A GPU pass that spins on a CPU-stream fence (MLX_METAL_FAST_SYNCH=1) had
no dependency on earlier passes still in flight on its queue. Without
hazard tracking Metal may start it first; when an earlier pass is what the
awaited CPU work needs (GPU -> CPU -> GPU, e.g. the sealed all-sum), the
GPU spins forever at 100 %. Traced on two M5 Ultra (GLM-5.3-Flash TP2): a
committed buffer of plain elementwise ops stayed scheduled-not-started
while the spin ran. The spin's pass now waits on every pending pass
(CommandEncoder::wait_for_all_pending); the CPU work it waits for already
follows them, so nothing is lost. Sealed TP2 then runs: prefill 1,913
tok/s, decode 50.7 tok/s (event sync instead: 13 tok/s).

The sealed all-sum's rank-order add is split across cores (256K-element
spans via dispatch_apply); the order of the two addends is unchanged, so
the result is bit-identical. 32 MiB fp32 sealed sum: 20-25 ms -> 5.9 ms.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…mand buffer

Port of the owner's mlx-m5u 8c4bf63fe: instead of one completion handler
per primitive retaining its inputs, a command buffer collects them and one
handler releases them all. Host time only; no kernel or graph changes.

GLM-5.3-Flash TP2 decode on two M5 Ultra (with the larger buffer-size
bound): 16.75 -> 15.92 ms a step together with the next commit; tokens
identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Port of the owner's mlx-m5u 42f2dcefb: a custom kernel's generated Metal
source and library name are built once per kernel and template instance
instead of on every launch. Host time only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rank JACCL exchange

On a two-rank JACCL mesh, `SealedAllSum` now runs one CPU-stream task in
which each 256 KiB sealed piece (its own AES-256-GCM tag) is sent as soon as
it is sealed and is opened and added as soon as it lands, so only the first
and last pieces wait for crypto. Other groups (ring, TCP, more ranks) keep
the v1 seal-gather-open path unchanged.

- `MeshImpl::exchange` (new, `mesh_impl.h`): receives posted before any send,
  at most two frames in flight each way, a frame leaves only when its bytes
  are sealed (`ready` is polled, never blocks, so a landed receive is always
  copied out and reposted); no self-copy, smallest size class for the last
  frame, frame schedule a function of n only.
- `sealed_stream.h`: piece layout (P = 256 KiB - 16, four sealed pieces per
  1 MiB frame), seal lanes capped at 8 on the user-interactive queue, the
  add in place per piece with v1's expression and rank-0-first operand
  order (bit-identical results).
- Security as v1: every piece is one AES-GCM box under its direction and
  generation key, nonce (piece, sender, counter) sealed once; the receiver
  derives piece count and lengths from n and P, never from the wire; the
  authenticated data is v1's plus P, so neither layout's piece verifies as
  the other's; a piece that does not open latches the codec and fails the
  sum, and a refused seal sends zeros so both ranks refuse in lockstep.
- v1 fix: a refused seal now zeroes its frame before throwing (the
  following gather could otherwise send a recycled buffer's stale bytes in
  the clear).
- C ABI: additive (`mlx_distributed_set_streamed_seal_codec`, two callback
  typedefs); `mlx_distributed_set_seal_codec` unchanged.

Measured on two M5 Ultra over Thunderbolt 5 (GLM-5.3-Flash TP2, sealed
bf16 partials, 15.7K-token prompt, 8,192-token chunks as two micro-batches):
prefill 3,554 -> 3,812 tok/s, output tokens identical. A tampered piece at
sealed counter 0 and at 200 is refused on both ranks, and both end without
hanging. `tests/jaccl/exchange_tests.cpp` (fake verbs, real MeshImpl and
streamed-sum code, stand-in keyed codec) passes under ASan/UBSan and TSan:
sizes 1 B to 3 MiB, bitwise equality with v1 for fp32/bf16/fp16 over 1-65
pieces, tamper, seal failure, closed group.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With `MLX_METAL_FUSED_FENCE_UPDATE=1` (and `MLX_METAL_FAST_SYNCH=1`), the
`input_coherent` rewrite, the system fence and `fence_update` become one
dispatch (`fence_update_coherent`): every word of the partial still gets a
coherent store and a system fence before the timestamp; larger inputs use a
monotonic arrival count. Off by default; with it off the dispatch sequence
is today's. Synchronization only: outputs are unchanged (GLM-5.3-Flash TP2
on two M5 Ultra: 256 tokens identical, no hangs).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With `MLX_METAL_FENCED_INPUTS=1` (needs fast synch), a custom kernel can
declare an input produced by a CPU-stream collective as fenced: thread 0
of each threadgroup waits on the fence before reading it (`MLX_FENCED_INPUT_WAIT`),
and MLX skips the separate `fence_wait` dispatch. The pass still waits on
every pending pass first (`wait_for_all_pending`, the ordering that fixed
the fast-synch deadlock), plus a barrier before the kernel. Falls back to
`fence_wait` when the input needs a copy, another GPU stream feeds the CPU
in the same eval, or the graph is exported. Off by default; outputs are
unchanged (tokens identical on the GLM TP2 pair).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`CommandEncoder` tracked each dispatch's inputs and outputs in three hash
sets; decode-sized dispatches (<= 24 buffers) now use small flat vectors
with identical barrier, commit and fence decisions (replayed against the
old code on 100 decode-like traces; set differential test under ASan/UBSan).
About 0.2 ms of host time per GLM-5.3-Flash decode step (~1,000 dispatches).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`MLX_EAGER_FENCE_COMMIT=1` commits the command buffer right after a
GPU-to-CPU fence update, so a CPU-stream collective can start as soon as
its input is ready. Off by default (measured neutral on GLM TP2 decode);
it only adds commit points and keeps the fast-synch ordering intact.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`MLX_MAX_ACTIVE_CPU_TASKS=n` counts CPU-stream tasks apart from command
buffers in the active-task throttle (`MLX_MAX_ACTIVE_TASKS` sets the
other). Unset, the throttle is exactly today's.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Host time only; every generated source, kernel name, buffer and dispatch
order is unchanged:
- custom-kernel primitives share the cached generated source instead of
  copying it, and resolve the pipeline once per source;
- repeat launches skip rebuilding the kernel name, template string and
  cache key; inputs bind directly when none needs a contiguous copy;
- graph teardown dedups node inputs with a small vector; `eval_impl`'s
  per-node map is a flat open-addressing map (differential test vs
  std::unordered_map under ASan/UBSan);
- one per-buffer Metal query removed from allocation.

GLM-5.3-Flash one-rank decode on an M5 Ultra (with the mlx-c/mlx-swift
counterparts): main-thread CPU per step 6.5–6.9 -> 4.7–4.8 ms; logits and
token hashes identical. The CUDA side of the shared-source change is not
compiled here.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308 and others added 4 commits October 21, 2026 10:36
This reverts commit 4d91320, as the single-Mac core did (mlx#37 4a35574,
6ee3062). The one-page minimum charged MTP drafters without declared layouts
37-671 MB per token and made each GLM MTP request reserve about 11.3 GB fixed
instead of 4.1 GB. The single-Mac revert cannot reach this branch through the
merge (its net change there is zero), so it is applied here explicitly.
unified_memory.rst and process_allocation_cap_metal.cpp now match the
single-Mac core; allocator.cpp keeps this branch's cached maxBufferLength.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The fork-diff gate failed on this branch since 7620956 and 8c6f9e7:
mlx/array.cpp and mlx/distributed/sealed_stream.h were changed but not
described. Add a section for the custom-kernel launch and graph bookkeeping
cuts, and the streamed pieces to the sealed all-sum section.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
origin/main is the squash merge of mlx#37, whose head ad59668 this branch
already merged, so the tree is unchanged; this records main as merged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`reduce_scatter_sealed` and `all_gather_sealed` (SealedReduceScatter,
SealedAllGather) seal and authenticate every piece through the same codec
as `all_sum_sealed`. The streamed path generalises the streamed all-sum
(source and addend sealed apart; the gather adds nothing), with a
gathered path for groups without exchange. The reduce-scatter adds each
element's two partials in float32 in rank order and rounds once, as the
all-sum does; the gather is a byte copy.

Every sealed collective now draws its counter from one process-wide
sequence (`next_sealed_counter`, which was the all-sum's own counter):
the same values for a graph of all-sums only, and no counter is reused
across collective kinds. A peer piece that does not authenticate is
refused, and the codec stays refused.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308 and others added 3 commits October 25, 2026 23:08
…ectives branch

Brings in merge/33 at 187b6c6, which is main f78f0fa (the upstream
ml-explore/mlx sync to afc03c5 with its restored history, and the fork
fixes #40, #43-#46) merged into the JACCL send-frame and progress-guard
branch. This branch had merged main up to ad4745a; everything main added
since comes in here.

Conflicts and how they were resolved:
- mlx/fence.h, mlx/backend/{metal,cuda,no_gpu}/fence.cpp, mlx/transforms.cpp:
  main's value-based fence API (upstream ml-explore#4552: update returns a value,
  wait takes it) with this branch's fenced inputs and eager commit.
  Fence::wait_in_consumer now takes the value of the update it waits for,
  and FencedInput carries that value, so a fenced input never waits for a
  later update than its own. The Metal fence keeps the two-word buffer of
  the fused coherent update, now owned through main's allocator::Data.
- Metal fast fence wait: this branch orders the spin with
  wait_for_all_pending (the pass waits for every pass in flight). Upstream
  ml-explore#4689 gets the same order by putting the spin in a pass of its own
  (end_encoding() before it, end_encoding(true) after it). The merge keeps
  the spin in the current pass and adds a buffer barrier after
  wait_for_all_pending, so it also follows the earlier kernels of the same
  pass, as upstream's separate pass does (the same line that the next
  branch, fix/cluster-metal-fence-wait, adds). A wait adds no encoder
  boundary; tensor-parallel decode waits for a CPU-stream all-sum about
  twice a layer. end_encoding(bool wait_in_flight) stays and shares one
  loop with wait_for_all_pending.
- docs/src/usage/environment_variables.rst: both sides' entries
  (MLX_METAL_FUSED_FENCE_UPDATE, MLX_METAL_FENCED_INPUTS, MLX_CONV_WINOGRAD).
- fork.yaml: the merged JACCL branch's file plus this branch's sections
  (startup readiness, two-rank sealed all-sum, fused fence update, fenced
  inputs, cheaper custom-kernel launches). The fenced-input section now
  explains the fence-wait order against upstream ml-explore#4689 and the value-based
  waits. The "not fork work" backport section is gone (upstream ancestry).

Also adapted:
- mlx/distributed/jaccl/lib/jaccl/mesh.h: main's mesh has one MeshImpl per
  wire (upstream ml-explore#4607). The startup wait budget goes to every wire; the
  two-rank exchange of the streamed sealed sum runs on the first wire. The
  two-Mac cluster's device matrix has one device per peer, so it has one
  wire and nothing changes for it.

Unchanged by the merge: main's quantized, matmul and SDPA kernels and their
dispatch (this branch had no change there beyond what main already carried
since #37), main's revert of the upstream ml-explore#4352 simdgroup skip, and this
branch's revert of 4d91320 (main does not carry 4d91320; the footprint
policy keeps {page, page, 0, 0, page}).

Tests (Ultra A, M5 Ultra; Debug build with Metal, 8 jobs):
- tests/tests on the GPU: 310 of 316 cases passed. The 6 failures (linalg
  QR, inverse, Cholesky, pseudo-inverse, eigh, LU) fail the same way on
  main f78f0fa on this machine.
- tests/tests with DEVICE=cpu: 316 of 316 passed.
- ctest process_allocation_cap: 15 of 15 passed.
- check_forkdiff.py --upstream-ref upstream/main: 137 changed files, all
  described, base afc03c5 is the merge-base.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The same change as 66edc77 on fix/cluster-metal-fence-wait (applied here
with cherry-pick, so that merge stays clean): clang-format and cmake-format
on the six paths that failed this pull request's Check Lint job
(mesh_impl.h, sealed.cpp, sealed_stream.h, fast_primitives.h,
tests/jaccl/CMakeLists.txt, tests/jaccl/exchange_tests.cpp). Layout only;
no source body changes.

Tests: clang-format 21 (the 21.x line that .pre-commit-config.yaml pins):
no changes left in any .h or .cpp file of the tree.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
360d11b (sealed reduce-scatter and all-gather) added the `addend`
argument to detail::StreamedSum, but tests/jaccl/exchange_tests.cpp still
called the old constructor, so the tests/jaccl project did not compile
(CI does not build it). The all-sum sends and adds its whole partial, as
all_sum_sealed does (send = addend = partial).

Tests: tests/jaccl with ASan and UBSan: 14 of 14 passed (exchange,
startup_wait_budget, startup_wait_budget_api, startup_wait_budget_groups,
send_frame and the nine progress_guard modes); with TSan: 14 of 14 passed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308 added a commit to Layr-Labs/mlx-swift that referenced this pull request Oct 26, 2026
Restitches the nested pins after the merge of main 6959486:
- Source/Cmlx/mlx: 360d11b9 -> 2b5253ad. That is the merge of the refreshed
  mlx#33 into Layr-Labs/mlx#35, with the formatting of the inherited cluster
  sources and the streamed-sum test fix on top.
- Source/Cmlx/mlx-c: 8a50797b -> 73720efb, the merge of mlx-c main 3a06ddbd
  into Layr-Labs/mlx-c#15.

Regenerated as tools/update-mlx.sh does from those pins:
- The mirrored Source/Cmlx/include headers, the embedded JIT and CPU
  sources and the Xcode project already match the regeneration.
- The include-framework copies, which only the Xcode framework build uses,
  had never been regenerated for the cluster pins. They now carry:
  - from the C layer (mlx-c-distributed_group.h, mlx-c-fast.h): the sealed
    codec and collectives, the startup-wait budget, fenced kernels and
    prepared launches;
  - from the core (mlx-backend-metal-device.h, mlx-distributed-distributed.h,
    mlx-fast.h, mlx-scheduler.h, mlx-utils.h): the core's own header changes.
- Metal.hpp differs only in its "Autogenerated" line and is kept as it is.

fork.yaml's pin paragraph names the merged core and C.

Tests (Ultra A, Xcode 27.1, at these pins):
- tools/update-mlx.sh regeneration from the two pins matches the tree
  (Metal.hpp aside, as above).
- check-generated-metal-sources.py: pass.
- swift build --build-tests: complete.
- Fork Tests, all 3 shards of the CI plan, with an AOT mlx.metallib from
  this core:
  - XCTest: 852 listed, 848 passed, and the 2 permitted skips.
  - The fresh-process allocation-cap test and both multi-process ring
    suites passed.
  - Swift Testing: 749 of 749 passed.
  - test-shards.py verify: every test ran exactly once.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308 added a commit to Layr-Labs/d-inference that referenced this pull request Oct 26, 2026
…ation

libs/mlx 79d2ebe8 -> 2b5253ad (Layr-Labs/mlx#35 merged with main), libs/mlx-swift 59ff16af -> 3e1fc78e (#56 merged with main; nests core 2b5253ad and mlx-c 73720efb), libs/mlx-swift-lm 84112479 -> d456a05d (#308 merged with main aee6ebb6, pinning mlx-swift 3e1fc78e). Master f9d1af3's provider code needs main's SDK API, which only the main-merged stack has. The cluster packages' Package.resolved drop swift-argument-parser, which the merged SDK no longer pulls (swift package resolve); provider-swift's is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Guard Apple dispatch, QoS and Blocks code. Preserve the Apple path and use synchronous authenticated frame callbacks elsewhere. Add CPU coverage for framing, failures and callback ownership.

Validated both selected branches on macOS with C++20 and warnings as errors; actual hosted Linux execution remains pending.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
(cherry picked from commit 9ec023f)
jonathan308 added a commit that referenced this pull request Oct 26, 2026
#35 now carries 9ec023f (cherry-picked as a135f49) so its Linux jobs build; #41 already had it, so its tree does not change (fork.yaml: #41's).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… all-sum branch

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@jonathan308
jonathan308 marked this pull request as ready for review October 27, 2026 18:49
@jonathan308
jonathan308 requested a review from davidtai October 27, 2026 18:49
jonathan308 and others added 2 commits October 27, 2026 13:24
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…sealed frames over INT_MAX

- metal_kernel.cpp: a fenced custom-kernel input is always declared in the
  device address space. A 0-d or small input was `constant`, and the GPU may
  load a fixed-size constant buffer before the kernel runs, so before
  MLX_FENCED_INPUT_WAIT, which could read the CPU-stream value too early.
- custom_kernel.cpp: take the fenced waits at the start of eval_gpu, before
  the launch checks that can throw. A throw left the wait in the
  thread-local list, where a later kernel at the same address took it.
- sealed.cpp: frame lengths go through detail::frame_extent, which refuses a
  frame of more than INT_MAX bytes instead of wrapping the int array shape
  (a float32 partial of 1 GiB gathers more than 2 GiB). Test added.
- sealed.h: state the codec contract: the counter restarts at 0 in each
  process (keys must not be reused across processes), and the codec's ctx
  must outlive the graphs built with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@davidtai

Copy link
Copy Markdown

Review fixes in be4b55b:

  • Fenced custom-kernel inputs are always in the device address space. A 0-d or small input was constant, and the GPU may load a fixed-size constant buffer before MLX_FENCED_INPUT_WAIT runs.
  • CustomKernel::eval_gpu takes its fenced waits before any launch check can throw. Before, a throw left the wait in the thread-local list for a later kernel at the same address.
  • Sealed frame lengths go through detail::frame_extent, which refuses frames over INT_MAX bytes instead of wrapping the int shape. Unit test added.
  • sealed.h documents the codec contract: the counter restarts at 0 in each process, so keys must not be reused across processes, and the codec ctx must outlive the graphs built with it.

Not changed (design question): if one rank fails before it takes its counter and the other rank does not, the counters desync and every later sealed op fails authentication. This fails closed, but nothing recovers it.

…d all-sum branch

mesh_impl.h keeps this branch's per-wire StartupWaitBudget and takes #33's
shared failure flag (std::shared_ptr), so a failed wire still closes every
wire of the group.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai added a commit that referenced this pull request Oct 27, 2026
… branch

Brings #35's fenced-input, fenced-wait and INT_MAX sealed-frame fixes and
#33's JACCL fixes under this branch's fence-order fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai and others added 2 commits October 27, 2026 18:01
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai added a commit that referenced this pull request Oct 27, 2026
…ering branch

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@davidtai
davidtai changed the base branch from darkbloom/jaccl-send-frame-progress-guard to main October 27, 2026 23:18
progress_guard_checks: the limit of the checks mode goes from 250 ms to
2000 ms. Every case is written in terms of the limit, so the ratio stays
the same: below-limit gaps are limit / 2 (ratio 1:2) and must not trip,
and over-limit stalls must still trip inside [limit, limit + slack]. At
250 ms a hosted runner that oversleeps a 125 ms gap by more than 125 ms
let the waiting rank time out and close its endpoint, so the sender then
failed with "mesh send: no completion". At 2000 ms the gap can absorb
1 s of delay.

startup_wait_budget_groups: the expiry cases set the fake fabric's idle
wait to 0. The guard samples the startup deadline once in 1024 empty
polls, and each empty poll slept up to 1 ms. Hosted runners stretch that
sleep to 4-8 ms, so 1024 polls passed the fake poll's 8 s deadline,
refuses() saw "The fake poll reached its deadline" instead of "absolute
startup", and the uncaught assertion aborted the process at about 8.1 s.
Spinning polls are what the hardware does. No library code changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai added a commit that referenced this pull request Oct 27, 2026
…g branch

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai added a commit to Layr-Labs/d-inference that referenced this pull request Oct 27, 2026
The closure carries the runtime fixes of the MLX review: the JACCL mesh failure flag
and 64-bit all_gather offsets (Layr-Labs/mlx#33), fenced inputs kept on device and sealed
frames over INT_MAX refused (Layr-Labs/mlx#35), the fast fence before the event fence
(Layr-Labs/mlx#41), and GLM5NextClusterImagePlan.inspectPNG offsets taken from the Data's
startIndex (Layr-Labs/mlx-swift-lm#308). Pins: libs/mlx-swift-lm d5bd32b8, libs/mlx-swift c633d221 (the revision mlx-swift-lm's Package.swift pins), libs/mlx 72dffa2d (mlx-swift's Source/Cmlx/mlx). mlx-c in that mlx-swift is 2ac897e3.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@davidtai
davidtai merged commit 62829a8 into main Oct 28, 2026
26 checks passed
@davidtai
davidtai deleted the darkbloom/sealed-all-sum branch October 28, 2026 01:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants