Disaggregated Inference¶

Table of Contents¶

• Concept
• Why Separate Prefill from Decode?
• Full Disaggregated Flow
• Prefill Scheduler
• Decode Scheduler
• Request Router
• Stream Bridge
• KV Cache Transfer
• Serving Modes
• Configuration

Concept¶

Disaggregated Inference separates the two steps of LLM token generation onto different nodes:

1. Prefill (prompt processing): The entire input prompt is processed in a single forward pass to produce the initial KV cache and the first output token. This step is compute-bound: it requires many multiply-accumulate operations proportional to the prompt length.

2. Decode (autoregressive generation): Starting from the KV cache produced by prefill, new tokens are generated one at a time. This step is memory-bandwidth-bound: at each step, the model reads the full KV cache from memory but performs only a small amount of computation (one token at a time).

These two steps have different hardware requirements, and running them on the same GPU forces a compromise that is suboptimal for both.

Why Separate Prefill from Decode?¶

graph TB
    subgraph "Compute-Bound (Prefill)"
        PF_A["Dense matmul over<br/>full prompt length<br/>(e.g., 2048 tokens)"]
        PF_B["Benefits from:<br/>high FLOPs, tensor cores"]
    end

    subgraph "Memory-Bandwidth-Bound (Decode)"
        DC_A["One token forward pass<br/>per step, reads full KV cache<br/>(growing with each step)"]
        DC_B["Benefits from:<br/>high memory bandwidth,<br/>large VRAM"]
    end

    PF_A --- PF_B
    DC_A --- DC_B

By placing prefill on compute-dense hardware and decode on memory-dense hardware, both steps operate at their optimal point simultaneously, and the two pools can be scaled independently.

Operational advantages:

• Long prompts on prefill nodes do not delay ongoing decode sessions.
• Decode nodes are never blocked by a slow prefill of a long document.
• Prefill node capacity and decode node capacity can be scaled independently in response to the actual workload mix.

Full Disaggregated Flow¶

sequenceDiagram
    participant C as Client
    participant R as Router
    participant PF as Prefill Node
    participant DC as Decode Node

    C->>R: POST /v1/chat/completions
    R->>PF: Route request to prefill node
    Note over PF: Run prefill forward pass<br/>over prompt tokens<br/>Build KV cache for all layers
    PF->>PF: Sample first token
    PF->>R: Return first token + signal handoff
    R->>C: Stream first token (SSE)
    PF->>DC: Transfer KV cache (serialized + optionally compressed)
    DC->>DC: Admit sequence to decode batch
    loop Token generation
        DC->>DC: Decode step (one token)
        DC->>R: Stream token
        R->>C: Forward token (SSE)
    end
    DC->>R: Finish reason (stop / length)
    R->>C: data: [DONE]

The client API is identical to single-node serving.

Prefill Scheduler¶

src/distributed/disaggregated/prefill_scheduler.rs

The prefill node runs a dedicated scheduler that:

• Maintains a priority queue of pending prefill requests, ordered by prompt length (shorter-first) for better throughput.
• Bounds concurrent prefill operations by a configurable max_concurrent_prefills to prevent OOM under load.
• Monitors memory utilization and rejects new requests when a memory_threshold (default: configurable) is exceeded.
• Optionally runs chunked prefill: long prompts are processed in chunks of chunk_size_tokens, allowing progressive KV cache transfer to the decode node to start before the full prompt is processed.
• Serializes the KV cache after prefill completes and hands it off to the decode node via the TransferConfig.

Key configuration fields:

Decode Scheduler¶

src/distributed/disaggregated/decode_scheduler.rs

The decode node runs a scheduler that:

• Receives KV caches from prefill nodes asynchronously (non-blocking to ongoing decode steps).
• Admits sequences to the active decode batch when memory permits.
• Runs autoregressive generation: one decode step per active sequence per iteration.
• Emits completion events when sequences reach EOS, max_tokens, or an error, so the request router can deliver the final response.
• Applies backpressure to prefill nodes when the ingestion queue is full.

Key configuration fields:

Request Router¶

src/distributed/disaggregated/request_router.rs

The router is the entry point for all incoming API requests in disaggregated mode. Both POST /v1/chat/completions and POST /v1/completions are routed through the same prefill/decode handoff; the router:

1. Selects the best available prefill node using a configurable load-balancing strategy.
2. Selects the best available decode node and ships its address in the decode_target field of the prefill request frame, so the prefill node hands the KV cache off to the router-chosen decode node rather than only its own static --decode-peers configuration.
3. Tracks the full request lifecycle across the prefill → decode handoff.
4. Applies backpressure (queue, reject, or HTTP 503) when nodes are at capacity.
5. On node failure, automatically re-routes affected requests to healthy peers.

Load Balancing Strategies¶

Prefill-node and decode-node selection both use round-robin today, since the router has no live per-node load telemetry beyond its own dispatch counters.

Backpressure Actions¶

When prefill or decode nodes are at capacity, the router can take one of:

• Queue: Hold the request in the router's pending queue.
• Reject: Return HTTP 503 immediately.

The choice depends on the BackpressureConfig and the current load level (Low, Moderate, High, Critical). Every request passes admission control before dispatch, so a node at capacity returns HTTP 503 rather than being sent a request it cannot serve.

Health Monitoring and Failover¶

A background health monitor TCP-probes every registered peer's serving address on an interval, so a dead decode node (which the router never contacts directly on the request path) is detected and skipped even though the router only learns about prefill failures from a failed send. A transport error while delivering a prefill frame marks that node Unreachable, re-routes its in-flight requests to a healthy peer, and excludes it from further routing decisions until the health monitor observes it respond again. A failed first-token send or KV handoff fails only the one affected request; the prefill and decode workers keep serving the rest of their load. When admission finds no prefill node available while some are marked Unreachable, it re-probes those nodes first (single-flight, at most once per 250 ms router-wide, 500 ms connect timeout) and admits the request if one accepts, so a peer that came up after the last monitor cycle is not refused with HTTP 503. A full queue never triggers a probe.

Observability: GET /router/stats¶

GET /router/stats is mounted on the router's own client-facing HTTP port, alongside /v1/chat/completions. It reports per-node prefill/decode dispatch counts, node health, and a router metrics snapshot. By default the response redacts each peer's raw host:port and reports only stable router-assigned node ids (prefill-<i> / decode-<i>), enough to read the load spread and confirm a failed node is marked unreachable without disclosing the cluster topology to a client that can reach the inference port. Set MLXCEL_ROUTER_STATS_VERBOSE=1 on a trusted-segment deployment to opt back into the full address view.

As defense-in-depth on the KV handoff itself, a prefill node can validate an incoming decode_target against MLXCEL_DECODE_ALLOWLIST (a comma-separated host:port list covering every router-selectable decode node) before connecting, so a forged frame from a breached segment cannot redirect a KV cache handoff off-cluster. Leave it unset only when the router's client port and the internal prefill/decode ports already sit on the same trusted network segment.

Stream Bridge¶

src/distributed/disaggregated/stream_bridge.rs

The stream bridge ensures that the SSE token stream delivered to the client is seamless across the prefill-to-decode node boundary.

flowchart LR
    SB["**StreamBridge**\n\nPrefill: first token\n\nDecode: generated tokens"]
    SB -- "SSE stream" --> Client

The bridge tracks four stream states:

The bridge guarantees: - No gap in the token stream between the first token (from prefill) and subsequent tokens (from decode). - No duplicate tokens in the handoff. - Proper DONE signaling after the decode node finishes.

In the role-based serving deployment, the decode node sends each generated token as its own result frame, so the client sees per-token SSE deltas across the handoff rather than a buffered tail. The router does not rely on the remote stream for termination: it enforces the request's token budget itself and closes the stream when that budget is reached.

Router output passes through the same StreamFilter as single-node serving. Model-internal markers (<think>, tool-call delimiters, stray turn tokens) never reach the client, and thinking content is delivered as delta.reasoning_content in streaming responses and reasoning_content in non-streaming ones, matching single-node behavior byte for byte. Tool-call assembly into tool_calls is not yet supported on the router path; the filter suppresses tool-call delimiter markup from the stream.

KV Cache Transfer¶

src/distributed/kv_cache_transfer/

The KV cache produced by the prefill node must be transferred to the decode node before decoding can start. This is the primary latency cost of disaggregation.

mlxcel supports multiple transfer strategies:

graph LR
    PFN[Prefill Node] -->|transfer| DCN[Decode Node]

    subgraph "Transfer Strategies"
        FS[Full: send entire<br/>cache in one message]
        ST[Streamed: layer-by-layer<br/>as each completes]
        LP[LayerParallel: multiple layers<br/>concurrently]
    end

Adaptive Strategy Selection¶

AdaptiveSelector (src/distributed/kv_cache_transfer/adaptive.rs) measures runtime bandwidth between nodes using BandwidthEstimator and automatically selects the optimal transfer configuration. This avoids manual tuning of the transfer strategy per deployment.

Compression¶

Optional LZ4 compression is applied to KV cache data before transmission. Compression reduces transfer time on bandwidth-constrained links at the cost of additional CPU time on both ends.

Quantization¶

On-the-fly quantization of KV cache tensors is available: (src/distributed/kv_cache_transfer/quantized.rs):

• float16 → int8: 2× compression ratio.
• float16 → int4: 4× compression ratio (with accuracy trade-off).

Quantization reduces transfer bandwidth requirements at the cost of slight accuracy degradation in decode-step attention. For most models, int8 quantization has negligible quality impact.

KV Cache Serialization¶

src/distributed/kv_cache_serde/ handles serialization and deserialization of SerializableCacheState (the full KV cache) and SerializableSamplingState (the current sampling state to resume generation on the decode node).

Serving Modes¶

Defined in src/distributed/disaggregated/serving.rs:

Hybrid mode is zero-overhead: the HybridModeGuard wrapper detects when the mode is Hybrid and returns immediately (or delegates to the local model provider) without touching any distributed data structures. It is a direct passthrough, not "disaggregated with localhost transfer."

Configuration¶

Cluster Secret and Transport Protection¶

The prefill, decode, and router processes share one cluster secret, generated once with openssl rand -hex 32 > cluster.secret (at least 16 bytes) and copied to every host. Every serving-role connection runs the cluster handshake (protocol version 5): both ends prove possession of the secret with HMAC-SHA256 over the transcript, the accept or at-capacity verdict is authenticated, and every frame afterwards is encrypted and authenticated with AES-256-GCM under keys derived for that connection. A process built before v0.8.0 cannot join a v0.8.0 cluster, so upgrade every node together. See Pipeline Parallelism for the secret precedence and frame details.

Example: router, prefill, and decode on three processes¶

# Decode node: receives KV handoffs at --serving-bind.
mlxcel-server --model ./models/qwen3-0.6b-4bit \
  --node-role decode \
  --cluster-secret-file ./cluster.secret \
  --serving-bind 192.168.1.21:9100

# Prefill node: listens for request frames and hands KV off to the decode node.
mlxcel-server --model ./models/qwen3-0.6b-4bit \
  --node-role prefill \
  --cluster-secret-file ./cluster.secret \
  --serving-bind 192.168.1.20:9100 \
  --decode-peers 192.168.1.21:9100

# Router: model-free HTTP front (tokenizer + chat template only).
mlxcel-server --model ./models/qwen3-0.6b-4bit \
  --port 8080 \
  --node-role router \
  --cluster-secret-file ./cluster.secret \
  --serving-bind 192.168.1.10:9100 \
  --prefill-peers 192.168.1.20:9100 \
  --decode-peers 192.168.1.21:9100 \
  --no-warmup

Clients point at the router's HTTP port and use the regular OpenAI-compatible API; the merged response is byte-identical to single-node serving of the same checkpoint. Every process passes the same --model path. The router uses it only to load the tokenizer and chat template, so it needs the checkpoint files on disk but allocates no weight memory; --no-warmup skips the model warmup pass it cannot run.

Two operational bounds apply on the handoff path. Prompts longer than the prefill chunk size are chunk-prefilled to completion before the KV handoff, so long prompts work the same as on a single node. The network intake rejects prompts above an admission cap of 1,048,576 tokens.