Architecture Documentation

This document provides a detailed overview of the Secure Model-Distributed LLM Inference architecture.

Overview
System Architecture
Security Design
Communication Protocol
Verification Mechanism
Node Types
Data Flow

Overview

This system implements secure distributed inference for large language models by:

Partitioning the model across multiple nodes
Encrypting all inter-node communication
Verifying computational integrity using cryptographic commitments
Using pipeline parallelism for efficient generation

Key Design Goals

Privacy: No single node has access to the complete model
Security: All communications are encrypted and authenticated
Integrity: Computations can be verified without full re-execution
Efficiency: Minimize overhead while maintaining security guarantees

System Architecture

┌─────────────────────────────────────────────────────────────────┐
│                         Client Application                       │
└────────────────────────────┬────────────────────────────────────┘
                             │
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│                        Starter Node (Node 0)                     │
├─────────────────────────────────────────────────────────────────┤
│  Components:                                                     │
│  • Token Embeddings (wte)                                       │
│  • Position Embeddings (wpe)                                    │
│  • Dropout Layer                                                │
│  • Transformer Blocks 0-1                                       │
│  • Layer Normalization (ln_f)                                   │
│  • Language Model Head (lm_head)                                │
│                                                                  │
│  Responsibilities:                                               │
│  • Accept prompts and tokenize                                  │
│  • Generate embeddings                                          │
│  • Coordinate generation pipeline                               │
│  • Sample tokens from logits                                    │
│  • Return final generated text                                  │
└────────────────────────────┬────────────────────────────────────┘
                             │ Encrypted Channel
                             │ (AES-256-GCM + HMAC)
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│                      Secondary Node 1 (Node 1)                   │
├─────────────────────────────────────────────────────────────────┤
│  Components:                                                     │
│  • Transformer Blocks 2-7                                       │
│                                                                  │
│  Responsibilities:                                               │
│  • Receive hidden states from Node 0                           │
│  • Process through local transformer blocks                     │
│  • Create cryptographic proof of computation                    │
│  • Send to next node                                            │
└────────────────────────────┬────────────────────────────────────┘
                             │ Encrypted Channel
                             │ (AES-256-GCM + HMAC)
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│                      Secondary Node 2 (Node 2)                   │
├─────────────────────────────────────────────────────────────────┤
│  Components:                                                     │
│  • Transformer Blocks 8-11                                      │
│                                                                  │
│  Responsibilities:                                               │
│  • Receive hidden states from Node 1                           │
│  • Process through local transformer blocks                     │
│  • Create cryptographic proof of computation                    │
│  • Send back to Starter Node                                    │
└────────────────────────────┬────────────────────────────────────┘
                             │ Encrypted Channel
                             │ (AES-256-GCM + HMAC)
                             ▼
                    Back to Starter Node
                   (Pipeline Loop Continues)

Security Design

Multi-Layer Security

┌─────────────────────────────────────────────────────────┐
│                  Application Layer                       │
│  • Model partitioning ensures privacy                   │
│  • No node has complete model                           │
└────────────────────────┬────────────────────────────────┘
                         │
┌────────────────────────▼────────────────────────────────┐
│              Verification Layer                          │
│  • Cryptographic commitments to weights                 │
│  • Merkle tree commitments to activations               │
│  • Challenge-response proofs                            │
│  • Probabilistic verification (10% by default)          │
└────────────────────────┬────────────────────────────────┘
                         │
┌────────────────────────▼────────────────────────────────┐
│             Encryption Layer                             │
│  • AES-256-GCM for data encryption                      │
│  • HMAC-SHA256 for authentication                       │
│  • Per-session key exchange                             │
└────────────────────────┬────────────────────────────────┘
                         │
┌────────────────────────▼────────────────────────────────┐
│             Transport Layer                              │
│  • TCP sockets for reliable delivery                    │
│  • Optional TLS 1.3 wrapper                             │
│  • Connection management                                 │
└─────────────────────────────────────────────────────────┘

Security Features

1. Weight Commitment

commitment = SHA256(concat(all_model_parameters))

Created once at node startup
Included in every proof
Verifies model integrity

2. Activation Commitment

merkle_root = MerkleTree(flatten(activation_tensor))

Created for each forward pass
Allows efficient verification
Supports spot-checking

3. Encryption

Message Structure:
[4 bytes: length][32 bytes: HMAC][N bytes: AES-GCM(data)]

AES-GCM provides:
- Confidentiality (encryption)
- Integrity (authentication tag)
- Protection against replay attacks (nonce)

Communication Protocol

Message Format

class Message:
    sample_id: int              # Which generation sample
    data: torch.Tensor          # Activation tensor
    metadata: Dict              # Additional info
    proof: Dict                 # Verification proof
        ├─ computation_type
        ├─ input_commitment
        ├─ output_commitment
        ├─ weight_commitment
        └─ challenges[]

Connection Establishment

1. Node B listens on port
2. Node A connects to Node B
3. Key Exchange:
   Node A → generates session_key
   Node A → sends session_key to Node B
   Both nodes use session_key for AES-GCM
4. Connection ready for secure messaging

Message Exchange Flow

Sender Side:
1. Serialize message (pickle)
2. Compute HMAC of serialized data
3. Encrypt data with AES-GCM
4. Send: [length][HMAC][encrypted_data]

Receiver Side:
1. Receive length (4 bytes)
2. Receive HMAC (32 bytes)
3. Receive encrypted data (length bytes)
4. Decrypt data
5. Verify HMAC
6. Deserialize message
7. Put in queue for processing

Verification Mechanism

Commitment Scheme

# Weight Commitment (one-time)
H(W) = SHA256(concat(all_weights))

# Activation Commitment (per forward pass)
M(A) = MerkleRoot(hash(a_1), hash(a_2), ..., hash(a_n))

# Proof Structure
Proof = {
    'input_commitment': M(A_input),
    'output_commitment': M(A_output),
    'weight_commitment': H(W),
    'challenges': [random_spot_checks]
}

Verification Process

1. Node receives activation with proof
2. If verification_rate > random():
   a. Compute commitment of received activation
   b. Compare with proof's commitment
   c. Verify weight commitment matches
   d. Check challenge responses
   e. Record result (success/failure)
3. Continue processing regardless

Why Probabilistic Verification?

Performance: Verifying 10% vs 100% saves ~90% overhead
Security: Still provides strong integrity guarantees
Tunable: Can adjust verification rate based on trust level

Node Types

Starter Node

Unique Responsibilities:

Manages tokenization and detokenization
Generates embeddings (token + position)
Applies final layer norm and language model head
Samples tokens from logits
Orchestrates the generation loop

Key Methods:

def generate(prompts, max_new_tokens, temperature):
    # Initialize samples
    # For each iteration:
    #   1. Process through local chunk
    #   2. Send to next node
    #   3. Receive from last node
    #   4. Sample next token
    #   5. Update sample state
    # Return decoded texts

Secondary Node

Responsibilities:

Receive activations from previous node
Verify incoming proofs (probabilistically)
Process through local transformer blocks
Create proof of computation
Send to next node

Key Methods:

def _processing_loop():
    while running:
        # Receive message
        # Verify proof (if selected)
        # Forward through local model
        # Create new proof
        # Send to next node

Data Flow

Single Token Generation

1. Starter Node:
   [tokens] → embedding → blocks[0-1] → hidden_states₁
   
2. Send hidden_states₁ to Node 1 (encrypted)

3. Node 1:
   hidden_states₁ → blocks[2-7] → hidden_states₂
   
4. Send hidden_states₂ to Node 2 (encrypted)

5. Node 2:
   hidden_states₂ → blocks[8-11] → hidden_states₃
   
6. Send hidden_states₃ back to Starter (encrypted)

7. Starter Node:
   hidden_states₃ → layer_norm → lm_head → logits
   logits → sample_token → next_token

Pipeline Parallelism

For generating N tokens for M samples:

Time     Node 0          Node 1          Node 2
----------------------------------------------------
t=0      Sample 0        -               -
t=1      Sample 1        Sample 0        -
t=2      Sample 2        Sample 1        Sample 0
t=3      Sample 0 (t2)   Sample 2        Sample 1
t=4      Sample 1 (t2)   Sample 0 (t2)   Sample 2
...

All nodes work simultaneously on different samples/tokens!

Memory Management

Per Node:
├─ Model Chunk: ~300-500 MB (for GPT-2)
├─ Message Queues: ~10-50 MB
├─ Activation Buffers: ~1-10 MB
└─ Proof Cache: ~1-5 MB

Total: ~500 MB - 1 GB per node

Performance Characteristics

Latency Breakdown

Single Token Generation (CPU):
├─ Computation: 200-300ms per node × 3 nodes = 600-900ms
├─ Communication: 10-50ms per hop × 2 hops = 20-100ms
├─ Encryption/Decryption: 1-5ms per message × 2 = 2-10ms
├─ Verification (when triggered): 5-20ms
└─ Total: ~650-1000ms per token

With Pipeline Parallelism (3 samples):
├─ Initial fill: 650-1000ms
├─ Steady state: 200-400ms per token (3× speedup!)

Throughput

Batch Size    Tokens/sec    Efficiency
-----------------------------------------
1 sample      1.0-1.5       Baseline
3 samples     2.5-4.5       ~3× better
5 samples     3.5-6.0       ~4-5× better

Security Overhead

Component          Overhead
---------------------------------
AES-GCM           ~2-5%
HMAC              ~1-2%
Verification      ~5-10% (at 10% rate)
---------------------------------
Total             ~8-17%

Extension Points

Adding New Node Types

class CustomNode(BaseNode):
    def __init__(self, ...):
        super().__init__(...)
        # Custom initialization
    
    def start(self):
        # Custom startup logic
        super().start()

Custom Verification

class CustomCommitment(CommitmentScheme):
    def create_proof(self, input_act, output_act):
        # Custom proof generation
        proof = super().create_proof(input_act, output_act)
        # Add custom fields
        return proof

Different Partitioning Strategies

class CustomPartitioner(ModelPartitioner):
    def partition(self, num_nodes):
        # Custom partitioning logic
        # Could be layer-wise, block-wise, or hybrid
        return chunks

Security Considerations

Threat Model

Protected Against:

✅ Eavesdropping (encryption)
✅ Data tampering (HMAC + verification)
✅ Model theft (no single node has full model)
✅ Computation fraud (cryptographic proofs)

Not Protected Against:

⚠️ Compromised nodes (assumes honest-but-curious)
⚠️ Timing attacks (not hardened against side channels)
⚠️ DoS attacks (no rate limiting implemented)

Production Recommendations

Use TLS 1.3 for transport security
Implement proper key management (not hardcoded secrets)
Add node authentication (mutual TLS, certificates)
Enable audit logging for all operations
Add rate limiting to prevent abuse
Use hardware acceleration (AES-NI, GPU for verification)
Implement graceful degradation for node failures

For implementation details, see the source code in src/.

Navigation Menu

Architecture Documentation

Table of Contents

Overview

Key Design Goals

System Architecture

Security Design

Multi-Layer Security

Security Features

1. Weight Commitment

2. Activation Commitment

3. Encryption

Communication Protocol

Message Format

Connection Establishment

Message Exchange Flow

Verification Mechanism

Commitment Scheme

Verification Process

Why Probabilistic Verification?

Node Types

Starter Node

Secondary Node

Data Flow

Single Token Generation

Pipeline Parallelism

Memory Management

Performance Characteristics

Latency Breakdown

Throughput

Security Overhead

Extension Points

Adding New Node Types

Custom Verification

Different Partitioning Strategies

Security Considerations

Threat Model

Production Recommendations

FilesExpand file tree

ARCHITECTURE.md

Latest commit

History

ARCHITECTURE.md

File metadata and controls

Architecture Documentation

Table of Contents

Overview

Key Design Goals

System Architecture

Security Design

Multi-Layer Security

Security Features

1. Weight Commitment

2. Activation Commitment

3. Encryption

Communication Protocol

Message Format

Connection Establishment

Message Exchange Flow

Verification Mechanism

Commitment Scheme

Verification Process

Why Probabilistic Verification?

Node Types

Starter Node

Secondary Node

Data Flow

Single Token Generation

Pipeline Parallelism

Memory Management

Performance Characteristics

Latency Breakdown

Throughput

Security Overhead

Extension Points

Adding New Node Types

Custom Verification

Different Partitioning Strategies

Security Considerations

Threat Model

Production Recommendations