# Disaggregated Inference — Design Document ## Overview Disaggregated inference (DI) separates the **prefill** phase (prompt processing) from the **decode** phase (token generation) into independent server instances. This enables: - **Independent scaling** — prefill and decode have different compute profiles; scaling them separately matches hardware to workload. - **Latency isolation** — decode instances run exclusively decode batches, eliminating inter-token latency (ITL) spikes caused by prefill interference. - **Asymmetric parallelism** — each role can use different TP/DP/EP configurations optimized for its workload. This document covers the DI architecture as implemented in upstream vLLM and the Neuron-specific extensions in vllm-neuron. --- ## Architecture The diagram below illustrates the default **read mode** (decode pulls KV), which is what the toy proxy server drives. ```text +----------+ (1) request +-----------------------------+ | Client | ---------------> | Proxy Server | | | <--------------- | (load balancer; CPU-only, | +----------+ (5) SSE | no Neuron cores) | +-----------------------------+ | | (2) dispatch | | (3) dispatch (max_tokens=1) | | (w/ kv_transfer_params) v v +----------------------+ +----------------------+ | Prefill Server | | Decode Server | | (kv_producer) | | (kv_consumer) | | KV cache in device | | generates ALL out- | | memory | | put tokens, incl. | | | | the first | +----------------------+ +----------------------+ ^ | | (3') NIXL RDMA READ: | (4) stream | decode pulls KV | tokens +------------------------------+ ``` Flow (read mode): (1) client → proxy. (2) proxy dispatches to prefill (`max_tokens=1`) and blocks; prefill produces KV and returns `kv_transfer_params`. (3) proxy dispatches to decode with those params; (3') the decode worker pulls the KV from the prefill server via a NIXL RDMA READ. (4)/(5) decode generates and streams tokens back through the proxy to the client. ### Components | Component | Role | vLLM Configuration | |-----------|------|-------------------| | Prefill server | Processes prompt, produces KV cache (the one token it samples is **discarded**) | `kv_role: "kv_producer"` | | Decode server | Pulls KV, generates **all output tokens including the first** | `kv_role: "kv_consumer"` | | Proxy server | Load balancer / router — relays SSE responses. Runs **CPU-only (no Neuron cores)** | Application-level (not vLLM) | ### Request Flow (read mode) 1. Client sends request to the proxy server. 2. Proxy **dispatches the request to a prefill server** with `max_tokens=1`, `stream=false`, and `do_remote_decode=true`, and **blocks** on the response. 3. Prefill server computes the forward pass, producing KV cache in its device memory, and returns `kv_transfer_params` (engine id, block ids, host/port). The single token it samples is **discarded** — it exists only to drive KV production. 4. Proxy dispatches the request to a decode server, attaching the prefill's `kv_transfer_params`. When the decode worker starts the request (`start_load_kv`), it issues a **NIXL RDMA READ** to pull the KV blocks from the prefill server's device memory; the request waits in `WAITING_FOR_REMOTE_KVS` until the read completes. 5. Decode server runs autoregressive generation on the transferred KV — **the first output token is generated by the decode server** — and streams tokens back through the proxy to the client. > In **write mode**, the prefill server pushes KV to decode and the proxy can dispatch both concurrently (lower TTFT); see the Transfer Modes section. The diagram shows read mode for clarity. --- ## Upstream vLLM Implementation ### KV Connector Framework Located in `vllm/distributed/kv_transfer/`, the connector framework provides: - **`KVConnectorBase_V1`** (`kv_connector/v1/base.py`) — abstract interface with scheduler-side and worker-side methods. - **Scheduler-side**: `get_num_new_matched_tokens()`, `build_connector_meta()`, `update_state_after_alloc()`, `request_finished()` - **Worker-side**: `start_load_kv()`, `save_kv_layer()`, `wait_for_save()`, `register_kv_caches()` ### Connector Implementation vLLM ships several KV connectors, but Neuron DI uses **`NixlConnector`**: | Connector | Transport | Use Case | |-----------|-----------|----------| | `NixlConnector` | RDMA (NIXL library) | Production — GPU & Neuron | ### Transfer Modes | Mode | Direction | Proxy Behavior | Latency Impact | |------|-----------|---------------|----------------| | **Read** (decode pulls) | D ← P | Serial: await prefill, then dispatch decode | Higher TTFT | | **Write** (prefill pushes) | P → D | Concurrent: both dispatched simultaneously | Lower TTFT | ### V1 Engine Integration Points ```python # Scheduler integration connector.get_num_new_matched_tokens(request, num_computed_tokens) connector.update_state_after_alloc(request, blocks) meta = connector.build_connector_meta(scheduler_output) # Worker integration (context manager in model runner) with kv_connector_context(scheduler_output): connector.start_load_kv(forward_context) # Decode: load KV model.forward(...) # Compute connector.save_kv_layer(layer, kv, meta) # Prefill: save KV connector.wait_for_save() # Prefill: block until done ``` ### Request States ```text WAITING → WAITING_FOR_REMOTE_KVS → WAITING → RUNNING → FINISHED ↑ (KV transfer complete) ``` --- ## vllm-neuron Implementation ### Integration Strategy Neuron DI uses the upstream **`NixlConnector`** for all topologies, including hybrid-TP. EFA/LIBFABRIC is also upstream: it is selected via the `kv_connector_extra_config` `{"backends": ["LIBFABRIC"]}` option. ### Neuron-Specific DI Features #### Hybrid TP (Asymmetric Parallelism) Different tensor parallelism on prefill vs decode: ```text Prefill: TP=4 (fewer cores, optimized for compute-heavy prompt processing) Decode: TP=8 (more cores, optimized for memory-bandwidth-bound generation) ``` The connector handles the head-count mismatch by computing `head_ratio` and `inverse_head_ratio` during the NIXL handshake. #### Component DP (Decode-Side Module Sharding) Component DP shards individual model components (attention, embedding, MLP, lm_head) across DP replicas differently: ```json { "neuron_config": { "attention_dp_size": 4, "embedding_dp_size": 2, "mlp_dp_size": 4, "lm_head_dp_size": 2 } } ``` - Only supported on the decode server (`kv_role: "kv_consumer"`) - Requires `--data-parallel-size >= 2` - Forces the MoE engine path (DPEngineCoreProc) to maintain cross-DP process groups #### Speculative Decoding with DI (Eagle3) Eagle3 speculative decoding works with DI: - Prefill server produces the target KV cache **and runs spec decoding to generate the draft KV cache**, so both are available to the decode side. - Decode server runs the Eagle3 draft model + verification with rejection sampling. - Tested with DP+EP topologies (e.g., `dp4tp8_dp4tp8ep32`). #### Bi-directional KV Transfer (D→P) For multi-turn conversations, the decode server sends KV cache **back** to the prefill server so the next turn avoids recomputation: ```bash export VLLM_NEURON_BIDIRECTIONAL_KV=1 ``` The connector/scheduler handles the D→P direction (`is_bidirectional_kv_xfer_enabled`): - `get_num_new_matched_tokens`: Reports `remote_block_ids` from D as externally cached tokens. - `update_state_after_alloc`: Populates `_reqs_need_recv` for the prefill worker to pull KV from decode. - `request_finished`: Returns `kv_transfer_params` with `remote_block_ids` for the proxy to forward to the next prefill request. Supports a `remote_pull_threshold` — minimum token count before pulling from remote vs recomputing locally. #### Streaming KV Transfer Params When the decode server streams its response (SSE), `kv_transfer_params` are emitted in the final streaming chunk; the proxy reads them to relay decode block IDs back to the prefill server (used by bidirectional KV). This is part of **upstream vLLM (≥ 0.21.0)**. #### Role-Specific Warmup Environment variables to skip compilation for the unused phase: ```bash VLLM_NEURON_SKIP_PREFILL_WARMUP=1 # Decode-only server VLLM_NEURON_SKIP_DECODE_WARMUP=1 # Prefill-only server ``` Saves compilation time — no need to compile prefill NEFFs on a decode-only server. --- ## Configuration Reference ### Prefill Server ```bash vllm serve $MODEL \ --tensor-parallel-size 8 \ --kv-transfer-config '{ "kv_connector": "NixlConnector", "kv_role": "kv_producer", "kv_buffer_device": "cuda", "kv_connector_extra_config": {"backends": ["LIBFABRIC"]} }' ``` ### Decode Server `--data-parallel-size` / `--enable-expert-parallel` are only needed for DP+EP (MoE) decode; a basic decode server omits them. ```bash vllm serve $MODEL \ --tensor-parallel-size 8 \ --kv-transfer-config '{ "kv_connector": "NixlConnector", "kv_role": "kv_consumer", "kv_buffer_device": "cuda", "kv_connector_extra_config": {"backends": ["LIBFABRIC"]} }' # DP+EP (MoE) decode adds: # --data-parallel-size 4 --enable-expert-parallel ``` ### Environment Variables | Variable | Default | Description | |----------|---------|-------------| | `VLLM_NIXL_SIDE_CHANNEL_HOST` | `127.0.0.1` | NIXL metadata exchange bind address | | `VLLM_NIXL_SIDE_CHANNEL_PORT` | `5557` | NIXL metadata exchange port | | `VLLM_NEURON_SKIP_PREFILL_WARMUP` | `0` | Skip prefill NEFF compilation (decode-only) | | `VLLM_NEURON_SKIP_DECODE_WARMUP` | `0` | Skip decode NEFF compilation (prefill-only) | | `VLLM_NEURON_BIDIRECTIONAL_KV` | `0` | Enable D→P KV transfer for multi-turn | | `NEURON_VISIBLE_DEVICES` | all | Core partitioning for single-node DI | ### Proxy Server The proxy is a CPU-only load balancer (no Neuron cores). A reference implementation and launch script are in `examples/vllm_neuron/vllm/disaggregated_inference/` — `toy_proxy_server.py` (round-robin over prefill/decode instances) and `server.sh`: ```bash python examples/vllm_neuron/vllm/disaggregated_inference/toy_proxy_server.py \ --prefiller-hosts --prefiller-ports 8100 \ --decoder-hosts --decoder-ports 8200 \ --port 8000 ``` --- ## Supported Topologies | Topology | Description | Single-Node | Multi-Node | |----------|-------------|:-----------:|:----------:| | xPyD | General multi-prefill multi-decode (covers 1P1D, 2P1D, 1P2D, …) | ✓¹ | ✓ | | Hybrid TP | Different TP per role | ✓ | ✓ | | DP+EP | Data parallel with expert parallel on decode | ✓ | ✓ | | Spec decode (Eagle3) + DI | Draft model on decode side | ✓ | ✓ | ¹ Single-node xPyD is supported as long as the device has enough cores and memory for all prefill + decode instances. --- ## Upstream vs. Plugin Responsibility | Capability | Upstream vLLM | vllm-neuron | |-----------|:------------:|:-----------:| | KV Connector framework (base interfaces) | ✓ | — | | NixlConnector (RDMA transfer) | ✓ | — | | LIBFABRIC/EFA backend | ✓ | — | | Hybrid TP head splitting | ✓ | — | | Bi-directional KV (D→P) | ✓ | ✓ (config/enablement) | | Streaming kv_transfer_params in SSE | ✓ (≥ 0.21.0) | — | | Component DP with DI | — | ✓ | | Role-specific warmup skip | — | ✓ | | Prefill/decode batch separation | — | ✓ (NeuronScheduler) | | Bucket padding for compiled NEFFs | — | ✓ (NeuronScheduler) | | Proxy server (toy) | ✓ (examples) | ✓ (examples) | | Production orchestration | — | AIBrix, production-stack | --- ## Data Flow (Detailed) ### Prefill Server (data flow) ```text Request arrives → NeuronScheduler → Pads to prefill bucket → Allocates KV blocks → NeuronModelRunner.execute_model() → _prepare_model_input() (input_ids, positions, slot_mapping) → with maybe_get_kv_connector_output(): → model.forward() [compiled NEFF] → connector.save_kv_layer() per attention layer → connector.wait_for_save() → sample one token (DISCARDED — only drives KV production; the first output token is produced by the decode server) → Return response (with kv_transfer_params for proxy) ``` ### Decode Server (data flow) ```text Request arrives (from proxy) → NeuronScheduler → Request enters WAITING_FOR_REMOTE_KVS state → connector.start_load_kv() triggers NIXL RDMA READ → Polls until transfer complete → Request moves to RUNNING → NeuronModelRunner.execute_model() → _prepare_model_input() (uses transferred KV blocks) → model.forward() [compiled decode NEFF] → sample token → Stream tokens back through proxy ``` ### NIXL Transfer ```text Decode worker Prefill worker ──────────── ───────────── 1. register_kv_caches() 1. register_kv_caches() → NIXL memory region registration → NIXL memory region registration → Create block descriptors → Create block descriptors 2. Handshake (one-time per P-D pair) ← Exchange: agent handles, block lens, base addrs 3. Transfer (per request) → Map block_ids to descriptor_ids → Issue RDMA READ (NIXL) → Device memory read by D → Poll for completion → Notify P to free blocks → Free KV blocks ``` --- ## Performance Characteristics ### Why DI Helps on Neuron 1. **Prefill and decode use different compiled NEFFs** — they can never share a batch anyway (unlike GPU which *can* mix). DI makes this explicit and eliminates scheduling contention. 2. **Dedicated decode instances maintain steady ITL** — no prefill compilation or execution interrupts decode. 3. **Asymmetric parallelism** exploits Neuron's core partitioning — prefill can use fewer cores with higher utilization, decode can use more cores with DP for throughput. ### Overheads | Overhead | Source | Mitigation | |----------|--------|-----------| | TTFT increase | KV transfer latency (RDMA) | Write mode overlaps transfer with prefill | | First-request latency | NIXL handshake (one-time) | Cached per P-D pair | | Memory | Decode holds KV for all in-flight requests | Tune `--max-num-seqs` | | Compilation time | Each role compiles only its buckets | `VLLM_NEURON_SKIP_PREFILL_WARMUP` / `VLLM_NEURON_SKIP_DECODE_WARMUP` | --- ## Future Directions - **Production orchestration** — AIBrix and production-stack integration for autoscaling, health checks, and SLO-aware routing.