feat(cluster): TP-default dispatch + TensorParallelEngine scaffold

provider-swift/Sources/ProviderCore/P2P/Parallelism.swift

@@ -0,0 +1,157 @@
+import Foundation
+
+// MARK: - Parallelism
+//
+// Selects the model-split strategy for a cluster session. Operator-visible
+// via the `--parallelism` flag on `darkbloom serve`. The default is `.auto`,
+// which picks `.tp` when both cluster peers can support tensor parallelism
+// for the loaded model, falls back to `.pp` when they can't, and falls back
+// to `.single` when no cluster peer is available.
+//
+// Why `.tp` wins by default on 2-Mac Thunderbolt 5: both Macs run all
+// transformer layers in parallel rather than taking turns, halving
+// single-stream decode latency. The per-layer allreduce cost is negligible
+// on TB5 (~80 Gbps, sub-ms per hidden-dim vector). PP stays as a fallback
+// for non-Llama models (until they get *TP variants) and for clusters
+// where the link bandwidth makes allreduces too expensive.
+
+public enum Parallelism: String, Sendable, CustomStringConvertible {
+    /// Tensor parallelism: both ranks run all layers in parallel; per-layer
+    /// allreduce synchronizes activations. Requires a *TP model variant
+    /// (e.g. `LlamaModelTP`).
+    case tp
+
+    /// Pipeline parallelism: rank 0 runs layers 0..N/2, rank 1 runs
+    /// layers N/2..N. One activation transfer per token. Works for any
+    /// model whose architecture exposes a layer-range entry point
+    /// (i.e. `callPartial` on Llama).
+    case pp
+
+    /// No cluster: rank 0 runs the full model alone. The fallback when no
+    /// peer is connected.
+    case single
+
+    /// Auto-select per `decide(...)`.
+    case auto
+
+    public var description: String { rawValue }
+}
+
+// MARK: - Decision
+
+extension Parallelism {
+
+    /// Inputs the dispatcher considers when `--parallelism auto` is set.
+    public struct DecisionInputs: Sendable {
+        public let operatorChoice: Parallelism
+        /// Number of peers in the cluster (including self). 1 → single-rank.
+        public let worldSize: Int
+        /// True iff the loaded model has a published TP variant (e.g. Llama
+        /// → true via `LlamaModelTP`). Other architectures fall back to PP
+        /// until they get their own *TP variant.
+        public let modelHasTPVariant: Bool
+        /// `attentionHeads` of the loaded model. Used to check divisibility.
+        public let attentionHeads: Int
+        /// `kvHeads` of the loaded model. Used to check divisibility.
+        public let kvHeads: Int
+        /// True iff the underlying RDMA / jaccl backend is initialized and
+        /// `DistributedGroup` can be created. False on macOS < 26.2, on
+        /// pre-M5 hardware, or when `rdma_ctl` reports disabled.
+        public let distributedGroupAvailable: Bool
+
+        public init(
+            operatorChoice: Parallelism,
+            worldSize: Int,
+            modelHasTPVariant: Bool,
+            attentionHeads: Int,
+            kvHeads: Int,
+            distributedGroupAvailable: Bool
+        ) {
+            self.operatorChoice = operatorChoice
+            self.worldSize = worldSize
+            self.modelHasTPVariant = modelHasTPVariant
+            self.attentionHeads = attentionHeads
+            self.kvHeads = kvHeads
+            self.distributedGroupAvailable = distributedGroupAvailable
+        }
+    }
+
+    /// Pick the actual parallelism strategy from the operator's choice and
+    /// the runtime capabilities. Honors explicit operator overrides as long
+    /// as they're achievable; falls back with a reason when they aren't.
+    ///
+    /// Returns the chosen strategy and a short, operator-readable reason for
+    /// the decision (suitable for logging at startup).
+    public static func decide(_ inputs: DecisionInputs) -> (Parallelism, reason: String) {
+        // worldSize=1 short-circuits everything: no cluster peer, no
+        // parallelism. Operator can pass `--parallelism tp` and we'll still
+        // end up here because there's no peer to be tensor-parallel with.
+        if inputs.worldSize < 2 {
+            return (.single, reason: "no cluster peer connected (worldSize=\(inputs.worldSize))")
+        }
+
+        switch inputs.operatorChoice {
+        case .single:
+            return (.single, reason: "operator selected --parallelism single")
+
+        case .pp:
+            return (.pp, reason: "operator selected --parallelism pp")
+
+        case .tp:
+            // Honor explicit TP if achievable. If not, fail closed rather than
+            // silently downgrade — operator asked for TP for a reason.
+            if !inputs.distributedGroupAvailable {
+                return (
+                    .single,
+                    reason: "operator selected --parallelism tp but DistributedGroup unavailable (RDMA / M5 capability missing) — refusing to silently downgrade to PP, falling back to single-rank"
+                )
+            }
+            if !inputs.modelHasTPVariant {
+                return (
+                    .single,
+                    reason: "operator selected --parallelism tp but the loaded model has no TP variant — refusing to silently downgrade to PP, falling back to single-rank"
+                )
+            }
+            if !canShard(heads: inputs.attentionHeads, worldSize: inputs.worldSize)
+                || !canShard(heads: inputs.kvHeads, worldSize: inputs.worldSize)
+            {
+                return (
+                    .single,
+                    reason: "operator selected --parallelism tp but model heads don't divide evenly across worldSize=\(inputs.worldSize)"
+                )
+            }
+            return (.tp, reason: "operator selected --parallelism tp; capabilities OK")
+
+        case .auto:
+            return autoDecide(inputs)
+        }
+    }
+
+    private static func autoDecide(_ inputs: DecisionInputs) -> (Parallelism, reason: String) {
+        // Auto: prefer TP if every capability lines up; else PP.
+        if !inputs.distributedGroupAvailable {
+            return (
+                .pp,
+                reason: "auto → pp: DistributedGroup unavailable (RDMA / M5 capability missing)"
+            )
+        }
+        if !inputs.modelHasTPVariant {
+            return (.pp, reason: "auto → pp: loaded model has no TP variant")
+        }
+        if !canShard(heads: inputs.attentionHeads, worldSize: inputs.worldSize)
+            || !canShard(heads: inputs.kvHeads, worldSize: inputs.worldSize)
+        {
+            return (
+                .pp,
+                reason:
+                    "auto → pp: model heads don't divide evenly across worldSize=\(inputs.worldSize) (attentionHeads=\(inputs.attentionHeads), kvHeads=\(inputs.kvHeads))"
+            )
+        }
+        return (.tp, reason: "auto → tp: all capabilities OK")
+    }
+
+    /// True iff `heads` is positive and divides evenly across `worldSize`.
+    public static func canShard(heads: Int, worldSize: Int) -> Bool {
+        worldSize > 0 && heads > 0 && heads % worldSize == 0
+    }
+}

provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift

@@ -0,0 +1,165 @@
+import CryptoKit
+import Foundation
+import MLX
+import MLXLLM
+import MLXLMCommon
+#if canImport(os)
+import os
+#endif
+
+// MARK: - TensorParallelInference
+//
+// TP path counterpart to EncryptedPipelineInference. Both engines run on
+// rank 0; their rank-1 counterparts (TensorParallelServer / EncryptedPipelineServer)
+// run on the other Mac. The dispatcher (Parallelism.decide) picks which pair
+// to instantiate at cluster-session-ready time.
+//
+// Key TP-vs-PP difference: in PP, only rank 0's first half of the model
+// runs at a time, then rank 1's second half runs, with an activation
+// transfer in between. In TP, BOTH ranks run all N layers in parallel,
+// with allreduce per layer (handled inside `LlamaModelTP`'s sharded
+// linear layers via the underlying `DistributedGroup`). One activation
+// transfer per token is replaced by 2N small allreduces per token; the
+// per-Mac compute drops by ~½ because each Mac runs half the heads.
+//
+// On Thunderbolt 5 the allreduce cost is sub-ms and the latency win is
+// roughly 2× over PP for single-stream decode.
+//
+// STATUS: scaffold only. The cluster runtime is currently a stub
+// (see ClusterCommand.swift line 431 — "Integrate ... with the coordinator
+// request queue"). Wiring the engine into the live inference loop, jaccl
+// DistributedGroup initialization, and the rank-1 serve loop are tracked
+// for a follow-up. This file defines the API shape so the dispatcher and
+// tests have something concrete to target.
+
+// MARK: - TensorParallelConfig
+
+public struct TensorParallelConfig: Sendable {
+    public let numLayers: Int
+    public let hiddenDim: Int
+    public let vocabSize: Int
+    public let worldSize: Int
+
+    public init(numLayers: Int, hiddenDim: Int, vocabSize: Int, worldSize: Int) {
+        self.numLayers = numLayers
+        self.hiddenDim = hiddenDim
+        self.vocabSize = vocabSize
+        self.worldSize = worldSize
+    }
+}
+
+// MARK: - TensorParallelEngine (rank 0)
+
+/// Drives tensor-parallel inference from rank 0's perspective.
+///
+/// Both ranks load `LlamaModelTP` and call `model.callAsFunction(input)` in
+/// lockstep; the sharded linear layers internally allreduce activations via
+/// the `DistributedGroup`. Both ranks produce identical logits; rank 0
+/// samples the next token and streams it to the consumer, then signals
+/// rank 1 to advance.
+///
+/// On a singleton group (worldSize=1), this degenerates to ordinary
+/// single-rank inference — the sharded layers' allreduce is a no-op.
+///
+/// On link failure, each step is retried up to 3 times with a 3-second
+/// delay (mirrors EncryptedPipelineEngine semantics). After 3 consecutive
+/// failures, `ClusterError.serviceUnavailable` is thrown.
+public actor TensorParallelEngine {
+    private let config: TensorParallelConfig
+    /// Held as `any LLMModel` so callers can pass either the fp16 variant
+    /// (`LlamaModelTP`) or the quantized variant (`LlamaModelTPQ`) without
+    /// the engine needing to be templated on the concrete type. Callers
+    /// (typically the dispatcher in `ClusterDiscovery`) are responsible
+    /// for passing a TP-capable model — passing a non-TP `LlamaModel`
+    /// would type-check but produce single-rank semantics.
+    private let model: any LLMModel
+    private let tokenizer: any Tokenizer
+    private var cache: [any KVCache]
+    private let session: ClusterSession
+
+    private let logger = Logger(
+        subsystem: "io.darkbloom.provider", category: "TensorParallelEngine")
+
+    public init(
+        config: TensorParallelConfig,
+        model: any LLMModel,
+        tokenizer: any Tokenizer,
+        session: ClusterSession
+    ) {
+        self.config = config
+        self.model = model
+        self.tokenizer = tokenizer
+        self.session = session
+        self.cache = []
+        logger.info(
+            """
+            TensorParallelEngine initialized: layers=\(config.numLayers), \
+            hidden=\(config.hiddenDim), vocab=\(config.vocabSize), \
+            worldSize=\(config.worldSize)
+            """)
+    }
+
+    /// Reset the KV cache to a fresh state. Called between distinct requests.
+    public func resetCache() {
+        cache = model.newCache(parameters: nil)
+        logger.info("TP KV cache reset (per-rank shard of \(self.cache.count) layers)")
+    }
+
+    // The actual decode loop is intentionally NOT implemented here yet —
+    // it depends on the cluster runtime integration that's still a stub
+    // in ClusterCommand.swift. The shape will be:
+    //
+    //   func generate(prompt: [Int], maxTokens: Int) -> AsyncStream<Int> {
+    //       AsyncStream { continuation in
+    //           Task {
+    //               // 1. Send prompt token IDs to rank 1 over ThunderboltLink.
+    //               // 2. Both ranks call model.callAsFunction (synced via allreduce).
+    //               // 3. Rank 0 samples token from logits, yields to consumer.
+    //               // 4. Rank 0 signals rank 1 with the chosen token + "continue".
+    //               // 5. Repeat until EOS or maxTokens.
+    //           }
+    //       }
+    //   }
+}
+
+// MARK: - TensorParallelServer (rank 1)
+
+/// Drives the rank-1 side of tensor-parallel inference. Symmetric to rank 0
+/// — both ranks run `model.callAsFunction(input)` in lockstep with allreduce
+/// synchronization. Rank 1 doesn't sample; it just provides compute for the
+/// allreduces and waits for rank 0's "next input" / "session end" signals
+/// over the ThunderboltLink control channel.
+public actor TensorParallelServer {
+    private let config: TensorParallelConfig
+    /// See `TensorParallelEngine.model` — same polymorphism rationale.
+    private let model: any LLMModel
+    private let peer: ClusterPeer
+
+    private let logger = Logger(
+        subsystem: "io.darkbloom.provider", category: "TensorParallelServer")
+
+    public init(
+        config: TensorParallelConfig,
+        model: any LLMModel,
+        peer: ClusterPeer
+    ) {
+        self.config = config
+        self.model = model
+        self.peer = peer
+        logger.info(
+            "TensorParallelServer initialized: layers=\(config.numLayers), worldSize=\(config.worldSize)"
+        )
+    }
+
+    // Same stub status as TensorParallelEngine — the serve loop will be:
+    //
+    //   func serve() async throws {
+    //       try await peer.serve(modelState: ..., inferenceHandler: { conn, _, _ in
+    //           // Loop:
+    //           //   Receive input token IDs from rank 0.
+    //           //   model.callAsFunction(input, cache) — synced via allreduce.
+    //           //   Receive sampled token from rank 0; advance cache.
+    //           //   On sessionEnd, exit loop.
+    //       })
+    //   }
+}

provider-swift/Sources/darkbloom/StartCommand.swift

@@ -2,6 +2,8 @@ import Foundation
 import ArgumentParser
 import ProviderCore
 
+extension Parallelism: ExpressibleByArgument {}
+
 struct Start: AsyncParsableCommand {
     static let configuration = CommandConfiguration(
         abstract: "Start the provider as a background service.",
@@ -37,6 +39,12 @@ struct Start: AsyncParsableCommand {
     @Flag(help: "Register as RDMA-capable and auto-discover Thunderbolt-connected peers for pipeline inference. Requires a Secure Enclave and a logged-in account.")
     var rdmaEnabled = false
 
+    @Option(
+        help:
+            "Cluster parallelism strategy: auto (default — prefer TP, fall back to PP), tp (tensor-parallel — both Macs run all layers in parallel; requires Llama-class model + M5 + rdma_ctl enabled), pp (pipeline-parallel — split layers across Macs), single (no clustering)."
+    )
+    var parallelism: Parallelism = .auto
+
     mutating func run() async throws {
         // GPU is required. Reject CPU fallback up-front so we never
         // come up reporting healthy and then silently churn at 0.5 tok/s.
@@ -252,6 +260,7 @@ struct Start: AsyncParsableCommand {
                     )
                     Task { await discovery.start() }
                     print("RDMA cluster discovery enabled — watching for Thunderbolt peers")
+                    print("Cluster parallelism preference: \(parallelism) (resolved at session handshake)")
                 }
             } catch {
                 printError("Warning: --rdma-enabled requires Secure Enclave support: \(error.localizedDescription)")