From 3722c626ddb6a31a80bb2a92433b559ff1f40b69 Mon Sep 17 00:00:00 2001 From: anupsv <6407789+anupsv@users.noreply.github.com> Date: Wed, 20 May 2026 22:28:32 -0700 Subject: [PATCH 1/4] feat(cluster): TP-default dispatch + TensorParallelEngine scaffold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the Parallelism dispatcher that selects between tensor-parallel, pipeline-parallel, and single-rank inference based on operator choice and runtime capability. TP wins by default on 2-Mac Thunderbolt 5 clusters because both Macs run all transformer layers in parallel with per-layer allreduce, roughly halving single-stream decode latency vs PP (where one Mac waits while the other runs). New: - ProviderCore/P2P/Parallelism.swift: enum {tp, pp, single, auto} + decide() that takes DecisionInputs (operatorChoice, worldSize, modelHasTPVariant, attentionHeads, kvHeads, distributedGroupAvailable) and returns the chosen strategy with an operator-readable reason. - ProviderCore/P2P/TensorParallelInference.swift: TensorParallelEngine (rank 0) and TensorParallelServer (rank 1) scaffolds, parallel to EncryptedPipelineEngine/EncryptedPipelineServer. The decode loop is intentionally TODO — the cluster runtime is currently a stub at ClusterCommand.swift:431 ("Integrate ... with the coordinator request queue"). Wiring will land in a follow-up that touches both PP and TP engines together. - darkbloom CLI: --parallelism flag on `serve` (auto, tp, pp, single). Defaults to auto. Logged at startup as the operator's preference; the final strategy is resolved at session-handshake time when the peer's capabilities are known. Auto fallback (Parallelism.decide auto path): - worldSize < 2 → single (no peer) - DistributedGroup unavailable (non-M5, rdma_ctl disabled, macOS < 26.2) → pp (TB5 transfer is enough; jaccl isn't required for PP) - modelHasTPVariant == false (non-Llama until they get *TP variants) → pp via callPartial - heads % worldSize != 0 → pp - everything OK → tp Explicit --parallelism tp fails closed when capabilities are missing: falls back to single rather than silently downgrading to pp, so a misconfiguration is visible to the operator instead of masked. Submodule bumps: - libs/mlx-swift → a83e602 (Layr-Labs/mlx-swift#3, sharded linear primitives) - libs/mlx-swift-lm → b06fa03 (Layr-Labs/mlx-swift-lm#25, LlamaModelTP) Tests (Tests/ProviderCoreTests/ParallelismTests.swift): 14 tests covering the auto-decide table, single-rank short-circuit, explicit operator overrides, fail-closed semantics for --parallelism tp, and the canShard divisibility helper. All pass. Related: #193 (upstream mlx-swift distributed deviation tracker). --- libs/mlx-swift | 2 +- libs/mlx-swift-lm | 2 +- .../ProviderCore/P2P/Parallelism.swift | 157 +++++++++++++++++ .../P2P/TensorParallelInference.swift | 158 ++++++++++++++++++ .../Sources/darkbloom/StartCommand.swift | 9 + .../ProviderCoreTests/ParallelismTests.swift | 138 +++++++++++++++ 6 files changed, 464 insertions(+), 2 deletions(-) create mode 100644 provider-swift/Sources/ProviderCore/P2P/Parallelism.swift create mode 100644 provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift create mode 100644 provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift diff --git a/libs/mlx-swift b/libs/mlx-swift index fa6a4e8b4..a83e602e0 160000 --- a/libs/mlx-swift +++ b/libs/mlx-swift @@ -1 +1 @@ -Subproject commit fa6a4e8b4ed98dd08da7f0384f05003541523d87 +Subproject commit a83e602e0f431043b30fadbf0b510ad405e245e3 diff --git a/libs/mlx-swift-lm b/libs/mlx-swift-lm index c2fbbdc71..b06fa03ac 160000 --- a/libs/mlx-swift-lm +++ b/libs/mlx-swift-lm @@ -1 +1 @@ -Subproject commit c2fbbdc71283eece0a1dd2f3b7b73f08fb9f6c0c +Subproject commit b06fa03ac23a4b586ea597433492649bb3ba3437 diff --git a/provider-swift/Sources/ProviderCore/P2P/Parallelism.swift b/provider-swift/Sources/ProviderCore/P2P/Parallelism.swift new file mode 100644 index 000000000..5f3c53ae3 --- /dev/null +++ b/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 + } +} diff --git a/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift b/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift new file mode 100644 index 000000000..8e71af188 --- /dev/null +++ b/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift @@ -0,0 +1,158 @@ +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 + private let model: LlamaModelTP + 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: LlamaModelTP, + 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 { + // 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 + private let model: LlamaModelTP + private let peer: ClusterPeer + + private let logger = Logger( + subsystem: "io.darkbloom.provider", category: "TensorParallelServer") + + public init( + config: TensorParallelConfig, + model: LlamaModelTP, + 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. + // }) + // } +} diff --git a/provider-swift/Sources/darkbloom/StartCommand.swift b/provider-swift/Sources/darkbloom/StartCommand.swift index bd254d454..61252b933 100644 --- a/provider-swift/Sources/darkbloom/StartCommand.swift +++ b/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)") diff --git a/provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift b/provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift new file mode 100644 index 000000000..d8546c7b8 --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift @@ -0,0 +1,138 @@ +import Foundation +import Testing +@testable import ProviderCore + +// MARK: - Helpers + +private func defaultInputs( + operatorChoice: Parallelism = .auto, + worldSize: Int = 2, + modelHasTPVariant: Bool = true, + attentionHeads: Int = 32, + kvHeads: Int = 8, + distributedGroupAvailable: Bool = true +) -> Parallelism.DecisionInputs { + Parallelism.DecisionInputs( + operatorChoice: operatorChoice, + worldSize: worldSize, + modelHasTPVariant: modelHasTPVariant, + attentionHeads: attentionHeads, + kvHeads: kvHeads, + distributedGroupAvailable: distributedGroupAvailable + ) +} + +// MARK: - Single-rank short-circuit + +@Test func parallelismShortCircuitsToSingleWhenNoPeer() { + let (chosen, reason) = Parallelism.decide(defaultInputs(worldSize: 1)) + #expect(chosen == .single) + #expect(reason.contains("worldSize=1")) +} + +@Test func parallelismShortCircuitsToSingleEvenIfOperatorAsksForTP() { + // worldSize < 2 always wins; operator choice can't conjure a peer. + let (chosen, _) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, worldSize: 1)) + #expect(chosen == .single) +} + +// MARK: - Auto-decide path (the default) + +@Test func parallelismAutoPicksTPWhenAllCapabilitiesAlign() { + let (chosen, reason) = Parallelism.decide(defaultInputs()) + #expect(chosen == .tp) + #expect(reason.contains("auto")) +} + +@Test func parallelismAutoFallsBackToPPWhenDistributedGroupUnavailable() { + // Capability gate failed (non-M5, or rdma_ctl disabled). PP can still + // run because callPartial only needs ThunderboltLink, not jaccl. + let (chosen, reason) = Parallelism.decide( + defaultInputs(distributedGroupAvailable: false)) + #expect(chosen == .pp) + #expect(reason.contains("DistributedGroup unavailable")) +} + +@Test func parallelismAutoFallsBackToPPWhenModelHasNoTPVariant() { + // Non-Llama model (e.g. Mistral, Qwen) until they get their own *TP + // variants. PP works via callPartial. + let (chosen, reason) = Parallelism.decide( + defaultInputs(modelHasTPVariant: false)) + #expect(chosen == .pp) + #expect(reason.contains("no TP variant")) +} + +@Test func parallelismAutoFallsBackToPPWhenHeadsDontDivide() { + // Odd head count can't shard across worldSize=2. + let (chosen, reason) = Parallelism.decide( + defaultInputs(attentionHeads: 33)) + #expect(chosen == .pp) + #expect(reason.contains("divide evenly")) +} + +@Test func parallelismAutoFallsBackToPPWhenKVHeadsDontDivide() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(kvHeads: 7)) + #expect(chosen == .pp) + #expect(reason.contains("divide evenly")) +} + +// MARK: - Explicit operator overrides + +@Test func parallelismHonorsExplicitPP() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .pp)) + #expect(chosen == .pp) + #expect(reason.contains("operator selected")) +} + +@Test func parallelismHonorsExplicitSingle() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .single)) + #expect(chosen == .single) + #expect(reason.contains("operator selected")) +} + +@Test func parallelismHonorsExplicitTPWhenAchievable() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp)) + #expect(chosen == .tp) + #expect(reason.contains("operator selected --parallelism tp")) +} + +// MARK: - Explicit TP refuses to silently downgrade + +@Test func parallelismExplicitTPFailsClosedWhenDistributedGroupUnavailable() { + // Operator asked for TP. Capability gate failed. We refuse to silently + // give them PP — that would mask a misconfiguration. Fall back to single + // and surface the reason. + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, distributedGroupAvailable: false)) + #expect(chosen == .single) + #expect(reason.contains("refusing to silently downgrade")) +} + +@Test func parallelismExplicitTPFailsClosedWhenModelHasNoTPVariant() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, modelHasTPVariant: false)) + #expect(chosen == .single) + #expect(reason.contains("refusing to silently downgrade")) +} + +@Test func parallelismExplicitTPFailsClosedWhenHeadsDontDivide() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, attentionHeads: 33)) + #expect(chosen == .single) + #expect(reason.contains("divide evenly")) +} + +// MARK: - Divisibility helper + +@Test func parallelismCanShardRequiresEvenDivision() { + #expect(Parallelism.canShard(heads: 32, worldSize: 2) == true) + #expect(Parallelism.canShard(heads: 33, worldSize: 2) == false) + #expect(Parallelism.canShard(heads: 64, worldSize: 4) == true) + #expect(Parallelism.canShard(heads: 0, worldSize: 2) == false) + #expect(Parallelism.canShard(heads: 32, worldSize: 0) == false) +} From 784c52d5f0cb849489ab3542b748a16fe80a6ade Mon Sep 17 00:00:00 2001 From: anupsv <6407789+anupsv@users.noreply.github.com> Date: Wed, 20 May 2026 23:13:04 -0700 Subject: [PATCH 2/4] chore(submodule): bump libs/mlx-swift-lm to pick up LlamaModelTPQ Picks up the quantized TP variant + matrix-level sharding tests on Layr-Labs/mlx-swift-lm#25 (commit 56a5ca6). This makes makeLlamaTP() available for routing the dispatcher to the right TP variant when a 4-bit / 8-bit quantized checkpoint is loaded. --- libs/mlx-swift-lm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/mlx-swift-lm b/libs/mlx-swift-lm index b06fa03ac..56a5ca60b 160000 --- a/libs/mlx-swift-lm +++ b/libs/mlx-swift-lm @@ -1 +1 @@ -Subproject commit b06fa03ac23a4b586ea597433492649bb3ba3437 +Subproject commit 56a5ca60b775465ebe25c83a7be825f13318ff3c From 12d7320cf0d04a2646ca5363ebaee347281bc802 Mon Sep 17 00:00:00 2001 From: anupsv <6407789+anupsv@users.noreply.github.com> Date: Thu, 21 May 2026 00:13:59 -0700 Subject: [PATCH 3/4] chore(submodule): bump libs/mlx-swift-lm to pick up LoRA + test fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Picks up Layr-Labs/mlx-swift-lm#25 commit 735c28a — drops the misleading LoRA conformance on LlamaModelTP/Q (returns [] with a docstring pointing at LlamaModel for LoRA workflows) and renames the testLlamaTPRejectsNonDivisibleHeads test to match what it actually exercises. --- libs/mlx-swift-lm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/mlx-swift-lm b/libs/mlx-swift-lm index 56a5ca60b..735c28ac0 160000 --- a/libs/mlx-swift-lm +++ b/libs/mlx-swift-lm @@ -1 +1 @@ -Subproject commit 56a5ca60b775465ebe25c83a7be825f13318ff3c +Subproject commit 735c28ac0e5e915b3d0e09e470efbd72a8e3472e From 0f85115f6e2a7824aabba0ebe2c9c3d630290c1f Mon Sep 17 00:00:00 2001 From: anupsv <6407789+anupsv@users.noreply.github.com> Date: Thu, 21 May 2026 00:18:33 -0700 Subject: [PATCH 4/4] fix(tp-engine): accept any LLMModel so quantized variant is usable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TensorParallelEngine / TensorParallelServer held `model: LlamaModelTP` as a concrete type, which meant the dispatcher couldn't pass the quantized `LlamaModelTPQ` returned by `makeLlamaTP(args:quantization:group:)`. Without this fix, the entire quantized-TP path was dead code — type checking would reject the quantized model at construction time. Both fields now hold `any LLMModel`. The engine needs only the callAsFunction(_:cache:) and newCache(parameters:) operations from the LanguageModel protocol (which LLMModel refines), and both LlamaModelTP and LlamaModelTPQ conform. Callers (the dispatcher in ClusterDiscovery, once wired) are responsible for passing a TP-capable model — a non-TP `LlamaModel` would type-check but produce single-rank semantics. Document this explicitly in the property comment. --- .../P2P/TensorParallelInference.swift | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift b/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift index 8e71af188..4f89ce06b 100644 --- a/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift +++ b/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift @@ -66,7 +66,13 @@ public struct TensorParallelConfig: Sendable { /// failures, `ClusterError.serviceUnavailable` is thrown. public actor TensorParallelEngine { private let config: TensorParallelConfig - private let model: LlamaModelTP + /// 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 @@ -76,7 +82,7 @@ public actor TensorParallelEngine { public init( config: TensorParallelConfig, - model: LlamaModelTP, + model: any LLMModel, tokenizer: any Tokenizer, session: ClusterSession ) { @@ -125,7 +131,8 @@ public actor TensorParallelEngine { /// over the ThunderboltLink control channel. public actor TensorParallelServer { private let config: TensorParallelConfig - private let model: LlamaModelTP + /// See `TensorParallelEngine.model` — same polymorphism rationale. + private let model: any LLMModel private let peer: ClusterPeer private let logger = Logger( @@ -133,7 +140,7 @@ public actor TensorParallelServer { public init( config: TensorParallelConfig, - model: LlamaModelTP, + model: any LLMModel, peer: ClusterPeer ) { self.config = config