feat(cluster): PP decode loop#197
Open
anupsv wants to merge 1 commit into
Open
Conversation
Implements greedy pipeline-parallel inference over ThunderboltLink with AES-256-GCM sealed activation tensors. Rank 0 runs layers 0..splitLayer, seals the activation, and exchanges it with rank 1 (which runs splitLayer..N + norm + lm_head + argmax) via new ppActivation/ppToken/ ppSessionEnd frame types (0x0B–0x0D). Wires EncryptedPipelineEngine/ Server into ClusterDiscovery as the fallback when jaccl bootstrap fails. - ClusterControlMessage: add ppActivation, ppToken, ppSessionEnd with JSON-encoded Codable payloads (PPActivationPayload, PPTokenPayload, PPSessionEndPayload); update frame-encoding comment table - EncryptedPipelineInference: rewrite EncryptedPipelineEngine to use new frame types; introduce PPClusterSession protocol (extends ClusterSessionSendable with receiveInferenceFrame + currentSessionKey) for mock-injectable testing; rewrite EncryptedPipelineServer.handleRequest to dispatch on ppActivation/ppSessionEnd with per-request cache reset - ClusterModelLoader: add loadLlamaModel(modelDirectory:) for PP (no jaccl DistributedGroup required, uses plain LlamaModel + callPartial) - ClusterDiscovery: add _ppEngine/_ppServer; add tryBuildRank0PPEngine/ tryBuildRank1PPServer; currentPPEngine()/currentPPServer() accessors; jaccl bootstrap failure on either rank now falls back to PP instead of aborting; ClusterPeer.serve routes ppActivation/ppToken/ppSessionEnd through inferenceHandler - Tests: PipelineParallelDecodeTests.swift — 17 tests covering raw values, JSON round-trips, engine construction, generate() loop, frame sequence, EOS stopping, ppSessionEnd as final frame, sealed activation decryptability, and KV cache slicing for both ranks
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Benchmark ResultsRunner: 1-provider-streaming1 providers, 1 users, 30 requests, concurrency=5, streaming=true
Latency Decomposition
Assertion Report: FAIL
1-provider-non-streaming1 providers, 1 users, 20 requests, concurrency=5, streaming=false
Latency Decomposition
Assertion Report: FAIL
7-provider-multi-model7 providers, 5 users, 50 requests, concurrency=10, streaming=true
Latency Decomposition
Assertion Report: FAIL
3-provider-high-concurrency3 providers, 10 users, 60 requests, concurrency=20, streaming=true
Latency Decomposition
Assertion Report: FAIL
1-provider-queue-saturation1 providers, 10 users, 40 requests, concurrency=15, streaming=true
Latency Decomposition
Assertion Report: FAIL
3-provider-20-users3 providers, 20 users, 60 requests, concurrency=10, streaming=true
Latency Decomposition
Assertion Report: FAIL
1-provider-scaling1 providers, 5 users, 30 requests, concurrency=10, streaming=true
Latency Decomposition
Assertion Report: FAIL
3-provider-scaling3 providers, 5 users, 30 requests, concurrency=10, streaming=true
Latency Decomposition
Assertion Report: FAIL
5-provider-scaling5 providers, 5 users, 30 requests, concurrency=10, streaming=true
Latency Decomposition
Assertion Report: FAIL
3-provider-heavy-100conc-10kb3 providers, 20 users, 100 requests, concurrency=100, streaming=true
Latency Decomposition
Assertion Report: FAIL
|
6 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ppActivation/ppToken/ppSessionEndframe types (0x0B–0x0D) with JSON-encoded Codable payloads; distinct from the legacyinferenceStep/inferenceTokenraw-bytes protocolEncryptedPipelineEngine(rank 0) andEncryptedPipelineServer(rank 1) fully rewritten with the new frame types,PPClusterSessionprotocol for mock-injectable testing, 3-retry logic, and actor-safe KV cache slicingClusterModelLoader.loadLlamaModelfor single-rankLlamaModelload (no DistributedGroup);ClusterDiscoverywires PP engine/server as jaccl-bootstrap-failure fallback on both ranksPipelineParallelDecodeTests.swift; all TP tests still passArchitecture notes
Double encryption kept:
sealedActivationinppActivationis AES-GCM sealed byTensorCrypto(inner), then the whole frame is AES-GCM sealed again byClusterFrame.encode/decode(outer). This matches the design established byinferenceStep/inferenceTokenand is intentional (defence-in-depth).TensorCrypto's docstring does not flag the outer layer as redundant.KV cache slicing:
EncryptedPipelineEngineinit takesmodel.newCache(parameters: nil).prefix(splitLayer).EncryptedPipelineServertakes.suffix(numLayers - splitLayer). Both are verified in tests.PPClusterSessionprotocol: extendsClusterSessionSendable(from PR 4b) withreceiveInferenceFrame() async throws -> DataandcurrentSessionKey() async throws -> SymmetricKey.ClusterSessionconforms; tests useMockPPSession. The actor-isolation conformance mirrors the existing TP pattern exactly.ClusterModelLoader:
loadLlamaModel(modelDirectory:)added alongside the existingload(modelDirectory:)(forLlamaModelTP). No removal of TP loader; PP just doesn't need the jaccl group.Split point:
ClusterDiscoveryusesnumLayers / 2as the defaultsplitLayerwhen building PP engine/server from the jaccl-failure fallback path. This is informational-only for now; PR 4d can expose this as a config knob.Test plan
swift buildpasses clean (one pre-existing warning inClusterCommand.swift, not introduced by this PR)swift test --filter "PipelineParallelDecode"— 17/17 passswift test --filter "TensorParallelDecode"— 17/17 pass (no regression)Need help on this PR? Tag
@codesmithwith what you need. Autofix is disabled.