neuron-tai/.scratch/distributed-gguf-runtime/architecture.md

Distributed GGUF Runtime Architecture

Product Stance

The platform optimizes for access to high-quality models, not lowest latency. Latency is acceptable if the user can run models that are otherwise unavailable to them. The hard context limit for the first serious distributed runtime should be 128K tokens. Longer context usually means the product is compensating for missing task decomposition, retrieval, or workspace summarization.

Current State

The current node has two materially different inference paths:

• Full local PyTorch model: calls Hugging Face model.generate(), so Transformers owns autoregressive decode and local KV cache.
• Distributed PyTorch route: bypasses model.generate(), calls individual layers with use_cache=False, and recomputes the full growing sequence for every generated token.

Current distributed data flow:

client request
  -> head node formats prompt
  -> for each output token:
       head tokenizes full current text
       head runs early layers over all tokens
       head sends full activation [batch, sequence, hidden] to next node
       middle nodes run their layers over all tokens
       tail returns one decoded token string
       head appends token to text

This is correct for small demos but not viable for large models. For GLM-5.2, a single 128K seam activation is roughly:

128K tokens * hidden_size 6144 * 2 bytes ~= 1.5 GiB per hop

Sending that every output token is the bottleneck.

Target State

Target distributed data flow:

client request
  -> tracker selects route and pins session
  -> head node creates session_id
  -> prefill:
       prompt is chunked
       each shard computes its layer range
       each shard appends local KV/state for its own layers
       activations cross only layer seams
  -> decode loop:
       head sends one new token / one-step hidden state
       each shard reads local KV/state for session_id
       each shard appends one step to local KV/state
       only one-step activation crosses seams
       tail returns logits/token

The KV cache remains local to the node that computed it. It is not sent to the next node and not read from a remote cache server during every decode step.

Client Feedback

Streaming responses are desirable when the backend and client transport support them. The product should stream token deltas when possible, and it must always provide realtime Generation Telemetry while the route is working.

The fallback behavior is a non-streaming final answer plus live telemetry. That fallback is acceptable for early route proofs or models/backends that cannot expose clean token deltas yet, but the preferred client experience is streamed output plus telemetry.

Minimum client-visible telemetry:

• route/session accepted
• selected model and quantization
• prefill phase started/completed
• decode phase started
• generated token count
• rolling tokens per second
• route health or retry/failure reason
• estimated billing units when available

Implementation options:

• Server-Sent Events or WebSocket for realtime progress
• polling endpoint for simple clients
• OpenAI-compatible streaming for clients that require token deltas

This means "no token streaming" is acceptable only as a fallback. "Silent wait for minutes" is not acceptable.

Artifact Plane

Artifact distribution is separate from execution.

model publisher
  -> produces model manifest
  -> creates GGUF / safetensors / tokenizer artifacts
  -> content-addresses every file/chunk
  -> publishes torrent/magnet + HTTP fallback metadata

node
  -> chooses model/layer range
  -> downloads needed files/chunks
  -> verifies hash
  -> advertises availability to tracker

Required manifest fields:

• model id and version
• upstream source repo and revision
• license
• architecture name
• tokenizer files and hashes
• quantization
• tensor-to-layer map
• file/chunk hashes
• optional GGUF split files
• supported runtime backends
• context cap
• KV/cache format descriptor

Execution Plane

The tracker selects routes using layer coverage and observed performance:

route = [
  head node: embeddings + layers 0..k
  middle nodes: contiguous layer ranges
  tail node: final layers + norm + lm_head
]

Route selection inputs:

• model id/version/quantization
• layer coverage
• node hardware
• measured prefill throughput
• measured decode throughput
• queue depth
• latency to neighboring nodes
• cache warmth for the requested prefix/session
• reliability/reputation

The route is sticky for the request/session. A new route means either a fresh prefill or restoring compatible KV snapshots.

KV Cache Ownership

KV/state ownership is by layer range:

session_id = request scoped id
node A owns layers 0..15 KV for session_id
node B owns layers 16..31 KV for session_id
node C owns layers 32..77 KV for session_id

The tracker does not own hot KV. It may know which nodes hold active KV for session accounting and failure handling.

Cache servers may store:

• prompt-prefix snapshots
• session checkpoints for retry
• cold reusable context blocks
• audit samples

Cache servers must not be in the per-token hot loop unless colocated with the compute node.

128K KV Budget

GLM-5.2 compressed DSA/MLA-style estimate from config:

layers = 78
kv_lora_rank = 512
qk_rope_head_dim = 64
dtype = bf16 = 2 bytes
context = 128K

per_token ~= 78 * (512 + 64) * 2 = 89,856 bytes ~= 87.75 KiB
128K total ~= 10.7 GiB
per layer ~= 137 MiB

This is feasible when sharded:

Layer count	Approx active KV at 128K
1	137 MiB
10	1.37 GiB
20	2.75 GiB
78	10.7 GiB

The exact runtime value depends on implementation and cache quantization, but the order of magnitude is acceptable.

Protocol Sketch

Prefill

POST /v1/sessions/{session_id}/prefill
Content-Type: application/octet-stream
X-Meshnet-Model: zai-org/GLM-5.2
X-Meshnet-Route-Id: ...
X-Meshnet-Token-Range: 0-2047
X-Meshnet-Shape: 1,2048,6144
X-Meshnet-Dtype: bfloat16

<activation bytes>

The receiver:

• validates route/session
• runs assigned layer range for that chunk
• appends local KV/state
• forwards resulting activation to next hop

Decode

POST /v1/sessions/{session_id}/decode-step
Content-Type: application/octet-stream
X-Meshnet-Model: zai-org/GLM-5.2
X-Meshnet-Position: 131072
X-Meshnet-Shape: 1,1,6144
X-Meshnet-Dtype: bfloat16

<one-step activation bytes>

The receiver:

• loads local KV/state by session_id
• runs one decode step for assigned layers
• appends one token position to local KV/state
• forwards one-step activation

GGUF / llama.cpp Integration

The target llama.cpp integration needs more than llama-server.

Required capabilities:

• load full GGUF locally for immediate single-node performance
• optionally load only selected tensors/layers
• execute a layer range against inbound hidden states
• expose outbound hidden states at a boundary
• own per-session KV/state for only the loaded layer range
• support prefill chunks and decode-step calls
• expose model-specific cache metadata for DSA/MLA without requiring the tracker to understand tensor internals

If llama.cpp cannot expose these as stable APIs today, the collaboration target is an upstream extension rather than a long-lived fork.

Failure Model

Alpha behavior:

• Route node drops during prefill: fail request and retry from scratch.
• Route node drops during decode: fail request unless a recent KV snapshot exists.
• Tracker restart: active sessions may be lost; completed billing records persist.
• Node restart: local hot KV is lost.

Later behavior:

• periodic KV snapshots for long sessions
• prefix cache reuse across requests
• route repair when a semantically equivalent node has the same model/layer range and compatible cache snapshot

Security And Trust

Activation/KV data can reveal user prompts. Public volunteer routes are not private. For sensitive workloads:

• use private swarms
• allow paid trusted nodes
• encrypt transport
• avoid storing hot KV on untrusted shared cache servers
• sample outputs for fraud/audit as already planned in alpha hardening