Skip to content

Pipeline-parallel inference across an attested Apple Silicon cluster (+ Gemma 4) - #473

Draft
crypt0fairy wants to merge 2 commits into
Layr-Labs:masterfrom
crypt0fairy:feat/cluster-pipeline-parallelism
Draft

Pipeline-parallel inference across an attested Apple Silicon cluster (+ Gemma 4)#473
crypt0fairy wants to merge 2 commits into
Layr-Labs:masterfrom
crypt0fairy:feat/cluster-pipeline-parallelism

Conversation

@crypt0fairy

@crypt0fairy crypt0fairy commented Jun 26, 2026

Copy link
Copy Markdown

Summary

Adds multi-device clustering so a co-located, individually-attested Apple
Silicon cluster serves as one logical provider, splitting a model layer-wise
across nodes over an encrypted MLX ring — without breaking Darkbloom's
operator-blind guarantee. This unlocks models too large for any single Mac
(e.g. Gemma-4-26B, GPT-OSS-120B) by pooling the cluster's combined unified memory.

Two commits:

  1. Pipeline-parallel cluster — the full clustering stack (trust, transport, decode, batching).
  2. Gemma 4 support — wires the Gemma 4 shard into the cluster.

Depends on the fork submodules: Layr-Labs/mlx#2 (ring backend), Layr-Labs/mlx-swift#8
(enable ring), Layr-Labs/mlx-swift-lm#53 (GPT-OSS/Gemma-4 pipeline shards).

Architecture

consumer ─(E2E)→ COORDINATOR ─(WS, head only)→ HEAD ─(encrypted ring)→ PEERS
                  sees ONE provider               serves + streams      own layer slice
  • Only the head connects to the coordinator; it registers as a single provider
    (the coordinator's Provider == one socket model is preserved). Peers join only
    the ring and loop in lockstep via a per-request control-round all_gather.
  • Layer split is memory-weighted (each node ∝ usable RAM), largest-remainder.
  • Activations are encrypted per hop: X25519 key agreement over the ring +
    per-token ChaCha20-Poly1305, AAD-bound to cluster/request/layer/seq. A tap on
    the wire sees only ciphertext, with forward secrecy.

Core modules (provider-swift/Sources/ProviderCore/Cluster/): ClusterRoster,
ClusterHandshake, ClusterLinkCrypto, LayerPartition, ClusterPlan,
MLXDistributed, ActivationCodec, ClusterPipeline, ClusterServer,
DistributedInferenceEngine, PipelineModelShard + GPT-OSS/Gemma-4 adapters,
and the ClusterBatch* continuous-batching path.

Trust & Attestation

The privacy guarantee rests on decrypting only inside attested hardware. The cluster
extends single-node attestation to every node:

  • Each member's Secure Enclave signs its attestation blob; the head relays them;
    the coordinator verifies each and sets cluster trust = min(member trust).
  • The coordinator signs a ClusterRoster; nodes verify it before forming pairwise
    ClusterHandshake links (so they only ring with coordinator-authorized peers).

Honest gaps (not yet closed — flagged for review):

  • Per-member attestation gather is opt-in (DARKBLOOM_CLUSTER_ATTEST=1), off by
    default; with it off the coordinator attests the head and peers are trusted
    transitively.
  • Relayed peers cap at self_signed (no per-peer MDM/MDA over their own link).
  • No per-member challenge/nonce freshness yet (member-blob replay window).
  • A dedicated coordinator-adjacent Cluster Attestation Service (per-node enroll +
    roster issuance) is the proposed path to close these without disturbing the
    one-provider-per-socket serving model.

Performance (measured, batch=1, Thunderbolt)

Clustering's value is fitting models that don't fit one node, not speeding up
ones that do (pipeline runs GPUs sequentially at batch=1 — see
docs/cluster-benchmark.md). Continuous batching recovers ~2.5× aggregate
throughput. TP/EP analysis + the TP comms floor are documented in
docs/architecture/cluster-tensor-expert-parallel.md.

Validation

  • Crypto/planning: ClusterCryptoTests (roster, handshake, link cipher, split).
  • Sharding correctness: monolithic-vs-shard logit oracles at 0.0000 diff for
    GPT-OSS and Gemma 4.
  • Real hardware: GPT-OSS-20B and Gemma-4-26B-A4B-it-qat-4bit each sharded
    across two Macs over Thunderbolt (encrypted ring, coherent generation, ~28 tok/s).

Robustness

An orphaned peer used to busy-spin its control loop on a dead ring (pegging a CPU
core indefinitely); it now backs off exponentially and exits.


View with Codesmith Autofix with Codesmith
Need help on this PR? Tag /codesmith with what you need. Autofix is disabled.

@vercel

vercel Bot commented Jun 26, 2026

Copy link
Copy Markdown

@crypt0fairy is attempting to deploy a commit to the EigenLabs Team on Vercel.

A member of the Team first needs to authorize it.

…ster

Adds multi-device clustering so a co-located, individually-attested Apple
Silicon cluster serves one logical provider, splitting a model layer-wise
across nodes over an encrypted MLX ring — without breaking Darkbloom's
operator-blind guarantee.

Core (provider-swift/Sources/ProviderCore/Cluster/):
- Trust: ClusterRoster (coordinator-signed membership, min-member trust),
  ClusterHandshake (SIGMA pairwise mutual auth over the roster),
  ClusterLinkCrypto (HKDF directional keys + per-token ChaCha20-Poly1305,
  AAD-bound to cluster/request/layer/seq).
- Planning: LayerPartition (memory-weighted largest-remainder split),
  ClusterPlan (topology from [cluster] config).
- Transport: MLXDistributed (Swift binding over mlx-c distributed collectives),
  MLXRingEnvironment, ActivationCodec (tensor<->sealed bytes).
- Decode: ClusterPipeline (the verified lockstep ring loop), ClusterServer
  (control-round request broadcast), DistributedInferenceEngine, the
  PipelineModelShard seam + the GPT-OSS shard adapter.
- Continuous batching: ClusterBatch{Scheduler,Server,Pipeline,Control}
  (~2.5x aggregate throughput, validated cross-Mac).
- Robustness: peer control-loop backs off and exits on a dead ring instead of
  busy-spinning (fixes an orphaned peer pegging a CPU core for hours).

Coordinator (Go): per-member cluster attestation relay + verify
(cluster trust = min member), cluster-provider registers as one mlx-swift
provider, dev local-cluster switches.

Executables: cluster-run, cluster-provider (+ --peer), and bench/smoke
harnesses (solo-bench, comms-bench, batch-*).

Docs: clustering.md, cluster-node-handshake.md, cluster-benchmark.md,
cluster-tensor-expert-parallel.md (incl. pipeline/TP/RDMA diagrams),
cluster-perf-backlog.md, jaccl-rdma-readiness.md.

Cluster model support: GPT-OSS (Gemma 4 added in the following commit).
Submodules point at the ring-enabled mlx-swift + pipeline-shard
mlx-swift-lm fork branches.
Wires the Gemma 4 pipeline shard (mlx-swift-lm) into the cluster so a Gemma 4
text model runs across the ring through the same engine/transport/crypto path
as GPT-OSS.

- Gemma4ShardAdapter: bridges Gemma4PipelineShard to PipelineModelShard.
- Dispatch: ClusterHeadBringup + cluster-run select the gemma4 shard by
  model_type; configInt now reads text-tower dims from the nested `text_config`
  (Gemma 4's multimodal config layout) with a top-level fallback.
- gemma4-shard-smoke: synthetic monolithic-vs-shard correctness oracle
  (0.000000 logit diff; covers MoE experts, k_eq_v full-attention, tied
  embeddings, final-logit softcap).

Validated end-to-end: gemma-4-26B-A4B-it-qat-4bit sharded across two Macs over
Thunderbolt (18/12 layer split, encrypted ring, coherent generation).
@crypt0fairy
crypt0fairy force-pushed the feat/cluster-pipeline-parallelism branch from 20ef807 to 900a1db Compare June 26, 2026 19:21
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.

1 participant